flip_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/impl/egblas/one_if_max_sub.hpp"

namespace etl {

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

    friend etl_traits<one_if_max_sub_transformer>;

    static constexpr bool gpu_computable = impl::egblas::has_sone_if_max_sub && all_row_major<T> && all_floating<T>;

private:
    sub_type sub; 

    std::vector<size_t> max_indices;

public:
    explicit one_if_max_sub_transformer(sub_type expr) : sub(expr), max_indices(etl::dim<0>(expr)) {
        for (size_t i = 0; i < etl::dim<0>(expr); ++i) {
            max_indices[i] = max_index(sub(i));
        }
    }

    value_type operator[](size_t i) const {
        size_t i_i = i / dim<1>(sub);
        size_t i_j = i % dim<1>(sub);
        return i_j == max_indices[i_i] ? value_type(1) : value_type(0);
    }

    value_type read_flat(size_t i) const {
        size_t i_i = i / dim<1>(sub);
        size_t i_j = i % dim<1>(sub);
        return i_j == max_indices[i_i] ? value_type(1) : value_type(0);
    }

    value_type operator()(size_t i, size_t j) const {
        return j == max_indices[i] ? value_type(1) : value_type(0);
    }

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

    template <typename Y>
    auto gpu_compute_hint(Y& y) const noexcept {
        decltype(auto) t1 = smart_gpu_compute_hint(sub, y);

        auto t2 = force_temporary_gpu_dim_only(t1);

        impl::egblas::one_if_max_sub(etl::dim<0>(y), etl::dim<1>(y), 1, t1.gpu_memory(), 1, t2.gpu_memory(), 1);

        t2.validate_gpu();
        t2.invalidate_cpu();

        return t2;
    }
    template <typename Y>
    Y& gpu_compute(Y& y) const noexcept {
        decltype(auto) t1 = select_smart_gpu_compute(sub, y);

        impl::egblas::one_if_max_sub(etl::dim<0>(y), etl::dim<1>(y), 1, t1.gpu_memory(), 1, y.gpu_memory(), 1);

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(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();
    }

    friend std::ostream& operator<<(std::ostream& os, const one_if_max_sub_transformer& transformer) {
        return os << "one_if_max_sub(" << transformer.sub << ")";
    }
};

template <typename T>
struct hflip_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend etl_traits<hflip_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit hflip_transformer(sub_type expr) : sub(expr) {}

    static constexpr bool matrix = is_2d<sub_type>; 

    value_type operator[](size_t i) const {
        if constexpr (matrix) {
            size_t i_i = i / dim<1>(sub);
            size_t i_j = i % dim<1>(sub);
            return sub[i_i * dim<1>(sub) + (dim<1>(sub) - 1 - i_j)];
        } else {
            return sub[etl::size(sub) - i - 1];
        }
    }

    value_type read_flat(size_t i) const {
        if constexpr (matrix) {
            size_t i_i = i / dim<1>(sub);
            size_t i_j = i % dim<1>(sub);
            return sub.read_flat(i_i * dim<1>(sub) + (dim<1>(sub) - 1 - i_j));
        } else {
            return sub.read_flat(etl::size(sub) - i - 1);
        }
    }

    value_type operator()(size_t i) const {
        return sub(etl::size(sub) - 1 - i);
    }

    value_type operator()(size_t i, size_t j) const {
        return sub(i, columns(sub) - 1 - j);
    }

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

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(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();
    }

    friend std::ostream& operator<<(std::ostream& os, const hflip_transformer& transformer) {
        return os << "hflip(" << transformer.sub << ")";
    }
};

template <typename T>
struct vflip_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend etl_traits<vflip_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit vflip_transformer(sub_type expr) : sub(expr) {}

    static constexpr bool matrix = is_2d<sub_type>; 

    value_type operator[](size_t i) const {
        if constexpr (matrix) {
            size_t i_i = i / dim<1>(sub);
            size_t i_j = i % dim<1>(sub);
            return sub[(dim<0>(sub) - 1 - i_i) * dim<1>(sub) + i_j];
        } else {
            return sub[i];
        }
    }

    value_type read_flat(size_t i) const {
        if constexpr (matrix) {
            size_t i_i = i / dim<1>(sub);
            size_t i_j = i % dim<1>(sub);
            return sub.read_flat((dim<0>(sub) - 1 - i_i) * dim<1>(sub) + i_j);
        } else {
            return sub.read_flat(i);
        }
    }

    value_type operator()(size_t i) const {
        return sub(i);
    }

    value_type operator()(size_t i, size_t j) const {
        return sub(rows(sub) - 1 - i, j);
    }

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

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(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();
    }

    friend std::ostream& operator<<(std::ostream& os, const vflip_transformer& transformer) {
        return os << "vflip(" << transformer.sub << ")";
    }
};

template <typename T>
struct fflip_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    friend etl_traits<fflip_transformer>;

    static constexpr bool gpu_computable = false;

private:
    sub_type sub; 

public:
    explicit fflip_transformer(sub_type expr) : sub(expr) {}

    value_type operator[](size_t i) const {
        if (dimensions(sub) == 1) {
            return sub[i];
        } else {
            return sub[etl::size(sub) - i - 1];
        }
    }

    value_type read_flat(size_t i) const {
        if (dimensions(sub) == 1) {
            return sub.read_flat(i);
        } else {
            return sub.read_flat(etl::size(sub) - i - 1);
        }
    }

    value_type operator()(size_t i) const {
        return sub(i);
    }

    value_type operator()(size_t i, size_t j) const {
        return sub(rows(sub) - 1 - i, columns(sub) - 1 - j);
    }

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

    // Internals

    template <typename V>
    void visit(V&& visitor) const {
        sub.visit(std::forward<V>(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();
    }

    friend std::ostream& operator<<(std::ostream& os, const fflip_transformer& transformer) {
        return os << "fflip(" << transformer.sub << ")";
    }
};

template <typename T>
requires(  cpp::specialization_of<etl::hflip_transformer, T>
            || cpp::specialization_of<etl::vflip_transformer, T>
            || cpp::specialization_of<etl::fflip_transformer, T>
            || cpp::specialization_of<etl::one_if_max_sub_transformer, T>)
struct etl_traits<T> {
    using expr_t     = T;                                  
    using sub_expr_t = std::decay_t<typename T::sub_type>; 
    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        = etl_traits<sub_expr_t>::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   = etl_traits<sub_expr_t>::is_temporary;  
    static constexpr bool  gpu_computable = T::gpu_computable;                     
    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 etl_traits<sub_expr_t>::size(v.sub);
    }

    static size_t dim(const expr_t& v, size_t d) {
        return etl_traits<sub_expr_t>::dim(v.sub, d);
    }

    static constexpr size_t size() {
        return etl_traits<sub_expr_t>::size();
    }

    template <size_t D>
    static constexpr size_t dim() {
        return etl_traits<sub_expr_t>::template dim<D>();
    }

    static constexpr size_t dimensions() {
        return etl_traits<sub_expr_t>::dimensions();
    }

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

} //end of namespace etl