pow.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/pow.hpp"
#include "etl/impl/egblas/pow_yx.hpp"

namespace etl {

template <typename T, typename E>
struct pow_binary_op {
    static constexpr bool linear      = true; 
    static constexpr bool thread_safe = true; 
    static constexpr bool desc_func   = true; 

    template <vector_mode_t V>
    static constexpr bool vectorizable =
        (V == vector_mode_t::SSE3 && is_single_precision_t<T>) || (V == vector_mode_t::AVX && is_single_precision_t<T>) || (intel_compiler && !is_complex_t<T>);

    template <typename L, typename R>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_spow_yx) || (is_double_precision_t<T> && impl::egblas::has_dpow_yx)
                                           || (is_complex_single_t<T> && impl::egblas::has_cpow_yx) || (is_complex_double_t<T> && impl::egblas::has_zpow_yx);

    static constexpr int complexity() {
        return 16;
    }

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

    static constexpr T apply(const T& x, E value) noexcept {
        return std::pow(x, value);
    }

    template <typename V = default_vec>
    static ETL_STRONG_INLINE(vec_type<V>) load(const vec_type<V>& x, const vec_type<V>& y) noexcept {
        // Use pow(x, y) = exp(y * log(x))
        auto t1 = V::log(x);
        auto t2 = V::mul(y, t1);
        return V::exp(t2);
    }

    template <typename X, typename Y, typename YY>
    static auto gpu_compute_hint([[maybe_unused]] const X& x, [[maybe_unused]] const Y& y, YY& yy) noexcept {
        decltype(auto) t1 = smart_gpu_compute_hint(x, yy);

        auto t2 = force_temporary_gpu(t1);

#ifdef ETL_CUDA
        T power_cpu(y.value);
        auto power_gpu = impl::cuda::cuda_allocate_only<T>(1);
        cuda_check(cudaMemcpy(power_gpu.get(), &power_cpu, 1 * sizeof(T), cudaMemcpyHostToDevice));

        T alpha(1.0);
        impl::egblas::pow_yx(etl::size(yy), alpha, power_gpu.get(), 0, t2.gpu_memory(), 1);
#endif

        return t2;
    }

    template <typename X, typename Y, typename YY>
    static YY& gpu_compute(const X& x, [[maybe_unused]] const Y& y, YY& yy) noexcept {
        smart_gpu_compute(x, yy);

#ifdef ETL_CUDA
        T power_cpu(y.value);
        auto power_gpu = impl::cuda::cuda_allocate_only<T>(1);
        cuda_check(cudaMemcpy(power_gpu.get(), &power_cpu, 1 * sizeof(T), cudaMemcpyHostToDevice));

        T alpha(1.0);
        impl::egblas::pow_yx(etl::size(yy), alpha, power_gpu.get(), 0, yy.gpu_memory(), 1);
#endif

        yy.validate_gpu();
        yy.invalidate_cpu();

        return yy;
    }

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

template <typename T, typename E>
struct precise_pow_binary_op {
    static constexpr bool linear      = true; 
    static constexpr bool thread_safe = true; 
    static constexpr bool desc_func   = true; 

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

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

    static constexpr int complexity() {
        return 16;
    }

    static constexpr T apply(const T& x, E value) noexcept {
        return std::pow(x, value);
    }

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

template <typename T, typename E>
struct integer_pow_binary_op {
    static constexpr bool linear      = true; 
    static constexpr bool thread_safe = true; 
    static constexpr bool desc_func   = true; 

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

    template <typename L, typename R>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_spow_yx) || (is_double_precision_t<T> && impl::egblas::has_dpow_yx)
                                           || (is_complex_single_t<T> && impl::egblas::has_cpow_yx) || (is_complex_double_t<T> && impl::egblas::has_zpow_yx);

    static constexpr int complexity() {
        return 16;
    }

    static constexpr T apply(const T& x, E value) noexcept {
        T r(1);

        for (size_t i = 0; i < value; ++i) {
            r *= x;
        }

        return r;
    }

    template <typename X, typename Y, typename YY>
    static auto gpu_compute_hint(const X& x, [[maybe_unused]] const Y& y, YY& yy) noexcept {
        decltype(auto) t1 = smart_gpu_compute_hint(x, yy);

        auto t2 = force_temporary_gpu(t1);

#ifdef ETL_CUDA
        T power_cpu(y.value);
        auto power_gpu = impl::cuda::cuda_allocate_only<T>(1);
        cuda_check(cudaMemcpy(power_gpu.get(), &power_cpu, 1 * sizeof(T), cudaMemcpyHostToDevice));

        T alpha(1.0);
        impl::egblas::pow_yx(etl::size(yy), alpha, power_gpu.get(), 0, t2.gpu_memory(), 1);
#endif

        return t2;
    }

    template <typename X, typename Y, typename YY>
    static YY& gpu_compute(const X& x, [[maybe_unused]] const Y& y, YY& yy) noexcept {
        smart_gpu_compute(x, yy);

#ifdef ETL_CUDA
        T power_cpu(y.value);
        auto power_gpu = impl::cuda::cuda_allocate_only<T>(1);
        cuda_check(cudaMemcpy(power_gpu.get(), &power_cpu, 1 * sizeof(T), cudaMemcpyHostToDevice));

        T alpha(1.0);
        impl::egblas::pow_yx(etl::size(yy), alpha, power_gpu.get(), 0, yy.gpu_memory(), 1);
#endif

        yy.validate_gpu();
        yy.invalidate_cpu();

        return yy;
    }

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

} //end of namespace etl