batch_softmax_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"

namespace etl {

template <typename A, bool Stable>
struct batch_softmax_expr : base_temporary_expr_un<batch_softmax_expr<A, Stable>, A> {
    using value_type = value_t<A>;                           
    using this_type  = batch_softmax_expr<A, Stable>;        
    using base_type  = base_temporary_expr_un<this_type, A>; 
    using sub_traits = decay_traits<A>;                      

    static constexpr auto storage_order = sub_traits::storage_order; 

    explicit batch_softmax_expr(A a) : base_type(a) {
        //Nothing else to init
    }

    template <same_dimensions_and_order<A> C>
    static void check([[maybe_unused]] const A& a, [[maybe_unused]] const C& c) {
        if constexpr (all_fast<A, C>) {
            static_assert(decay_traits<A>::size() == decay_traits<C>::size(), "Invalid size");
        } else {
            cpp_assert(etl::size(a) == etl::size(c), "Invalid size");
        }
    }

    // Assignment functions

    template <typename C>
    constexpr static batch_softmax_impl select_default_impl(bool no_gpu) {
        if (cudnn_enabled && all_homogeneous<A, C> && all_floating<A, C> && !no_gpu) {
            return batch_softmax_impl::CUDNN;
        }

        return batch_softmax_impl::STD;
    }

#ifdef ETL_MANUAL_SELECT

    template <typename C>
    static batch_softmax_impl select_impl() {
        return select_default_impl<C>(local_context().cpu);
    }

#else

    template <typename C>
    constexpr static batch_softmax_impl select_impl() {
        return select_default_impl<C>(false);
    }

#endif

    template <etl_expr C>
    void assign_to(C&& c) const {
        if constexpr (decay_traits<C>::storage_order == storage_order) {
            inc_counter("temp:assign");

            auto& a = this->a();

            standard_evaluator::pre_assign_rhs(a);

            check(a, c);

            constexpr_select auto impl = select_impl<C>();

            if constexpr_select (impl == batch_softmax_impl::CUDNN) {
                inc_counter("impl:cudnn");

                decltype(auto) a_gpu = smart_forward_gpu(a);

                if constexpr (Stable) {
                    impl::cudnn::stable_softmax(a_gpu, c);
                } else {
                    impl::cudnn::softmax(a_gpu, c);
                }
            } else if constexpr_select (impl == batch_softmax_impl::STD) {
                inc_counter("impl:std");

                if constexpr (Stable) {
                    for (size_t i = 0; i < etl::dim<0>(c); ++i) {
                        auto m = max(a(i));
                        c(i)   = exp(a(i) - m) / sum(exp(a(i) - m));
                    }
                } else {
                    for (size_t i = 0; i < etl::dim<0>(c); ++i) {
                        c(i) = exp(a(i)) / sum(exp(a(i)));
                    }
                }
            } else {
                cpp_unreachable("Invalid selection for batch_softmax");
            }
        } else {
            std_assign_evaluate(*this, c);
        }
    }

    template <typename L>
    void assign_add_to(L&& lhs) const {
        std_add_evaluate(*this, std::forward<L>(lhs));
    }

    template <typename L>
    void assign_sub_to(L&& lhs) const {
        std_sub_evaluate(*this, std::forward<L>(lhs));
    }

    template <typename L>
    void assign_mul_to(L&& lhs) const {
        std_mul_evaluate(*this, std::forward<L>(lhs));
    }

    template <typename L>
    void assign_div_to(L&& lhs) const {
        std_div_evaluate(*this, std::forward<L>(lhs));
    }

    template <typename L>
    void assign_mod_to(L&& lhs) const {
        std_mod_evaluate(*this, std::forward<L>(lhs));
    }

    friend std::ostream& operator<<(std::ostream& os, const batch_softmax_expr& expr) {
        return os << "batch_softmax(" << expr._a << ")";
    }
};

template <typename A, bool Stable>
struct etl_traits<etl::batch_softmax_expr<A, Stable>> {
    using expr_t     = etl::batch_softmax_expr<A, Stable>; 
    using sub_expr_t = std::decay_t<A>;                    
    using sub_traits = etl_traits<sub_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        = sub_traits::is_fast;                  
    static constexpr bool is_linear      = true;                                 
    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 = sub_traits::storage_order;            

    template <vector_mode_t V>
    using vectorizable = std::true_type;

    template <size_t DD>
    static constexpr size_t dim() {
        return decay_traits<A>::template dim<DD>();
    }

    static size_t dim(const expr_t& e, size_t d) {
        return etl::dim(e._a, d);
    }

    static size_t size(const expr_t& e) {
        return etl::size(e._a);
    }

    static constexpr size_t size() {
        return decay_traits<A>::size();
    }

    static constexpr size_t dimensions() {
        return decay_traits<A>::dimensions();
    }

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

} //end of namespace etl