clip.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/clip_value.hpp"

namespace etl {

template <typename T, typename S>
struct clip_scalar_op {
    template <typename V = default_vec>
    using vec_type = typename V::template vec_type<T>;

    static constexpr bool linear      = true; 
    static constexpr bool thread_safe = true; 

    template <vector_mode_t V>
    static constexpr bool vectorizable = intel_compiler && !is_complex_t<T>;

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_sclip_value)
                                           || (is_double_precision_t<T> && impl::egblas::has_dclip_value)
                                           || (is_complex_single_t<T> && impl::egblas::has_cclip_value)
                                           || (is_complex_double_t<T> && impl::egblas::has_zclip_value);

    static constexpr int complexity() {
        return 1;
    }

    S min; 
    S max; 

    clip_scalar_op(S min, S max) : min(min), max(max) {}

    constexpr T apply(const T& x) const noexcept {
        return std::min(std::max(x, min), max);
    }

#ifdef __INTEL_COMPILER

    template <typename V = default_vec>
    vec_type<V> load(const vec_type<V>& lhs) const noexcept {
        return V::min(V::max(lhs, V::set(min)), V::set(max));
    }

#endif

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

        auto t2 = force_temporary_gpu(t1);
        impl::egblas::clip_value(etl::size(y), T(1), min, max, t2.gpu_memory(), 1);
        return t2;
    }

    template <typename X, typename Y>
    Y& gpu_compute(const X& x, Y& y) const noexcept {
        smart_gpu_compute(x, y);

        impl::egblas::clip_value(etl::size(y), T(1), min, max, y.gpu_memory(), 1);

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

        return y;
    }

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

} //end of namespace etl