relu_derivative.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

namespace etl {

template <typename T>
struct relu_derivative_binary_op {
    static constexpr bool linear      = true;  
    static constexpr bool thread_safe = true;  
    static constexpr bool desc_func   = false; 

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

    template <typename L, typename R>
    static constexpr bool gpu_computable = cudnn_enabled;

    static constexpr int complexity() {
        return 1;
    }

    template <typename V = default_vec>
    using vec_type = typename V::template vec_type<T>;

    static constexpr T apply(const T& lhs, const T& rhs) noexcept {
        return lhs > 0.0 ? rhs : 0.0;
    }

    template <typename V = default_vec>
    static ETL_STRONG_INLINE(vec_type<V>) load(const vec_type<V>& lhs, const vec_type<V>& rhs) noexcept {
        auto t1 = V::round_up(V::min(V::set(T(1.0)), lhs));

        return V::mul(t1, rhs);
    }

    template <typename L, typename R, typename Y>
    static auto gpu_compute_hint(const L& lhs, const R& rhs, Y& y) noexcept {
        decltype(auto) t1 = smart_gpu_compute_hint(lhs, y);
        decltype(auto) t2 = smart_gpu_compute_hint(rhs, y);
        decltype(auto) t3 = force_temporary_gpu_dim_only(t2);

        impl::cudnn::relu_backward(t1, t2, t3);

        return t3;
    }

    template <typename L, typename R, typename Y>
    static Y& gpu_compute(const L& lhs, const R& rhs, Y& y) noexcept {
        decltype(auto) t1 = smart_gpu_compute_hint(lhs, y);
        decltype(auto) t2 = smart_gpu_compute_hint(rhs, y);

        impl::cudnn::relu_backward(t1, t2, y);

        return y;
    }

    static std::string desc() noexcept {
        return "relu_back";
    }
};

} //end of namespace etl