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

namespace etl {

template <typename T>
struct cbrt_unary_op {
    static constexpr bool linear      = true; 
    static constexpr bool thread_safe = true; 

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

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_scbrt) || (is_double_precision_t<T> && impl::egblas::has_dcbrt)
                                           || (is_complex_single_t<T> && impl::egblas::has_ccbrt) || (is_complex_double_t<T> && impl::egblas::has_zcbrt);

    static constexpr int complexity() {
        return 8;
    }

    static constexpr T apply(const T& x) {
        return std::cbrt(x);
    }

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

        auto t2 = force_temporary_gpu_dim_only(t1);

        T alpha(1.0);
        impl::egblas::cbrt(etl::size(y), alpha, t1.gpu_memory(), 1, t2.gpu_memory(), 1);

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

        T alpha(1.0);
        impl::egblas::cbrt(etl::size(y), alpha, t1.gpu_memory(), 1, y.gpu_memory(), 1);

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

        return y;
    }

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

template <typename TT>
struct cbrt_unary_op<std::complex<TT>> {
    using T = std::complex<TT>; 

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

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

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_scbrt) || (is_double_precision_t<T> && impl::egblas::has_dcbrt)
                                           || (is_complex_single_t<T> && impl::egblas::has_ccbrt) || (is_complex_double_t<T> && impl::egblas::has_zcbrt);

    static constexpr int complexity() {
        return 8;
    }

    static constexpr T apply(const T& z) {
        auto z_abs = std::abs(z);
        auto z_arg = std::arg(z);

        auto new_abs = std::cbrt(z_abs);
        auto new_arg = z_arg / 3.0f;

        return {new_abs * std::cos(new_arg), new_abs * std::sin(new_arg)};
    }

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

        auto t2 = force_temporary_gpu_dim_only(t1);

        T alpha(1.0);
        impl::egblas::cbrt(etl::size(y), alpha, t1.gpu_memory(), 1, t2.gpu_memory(), 1);

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

        T alpha(1.0);
        impl::egblas::cbrt(etl::size(y), alpha, t1.gpu_memory(), 1, y.gpu_memory(), 1);

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

        return y;
    }

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

template <typename TT>
struct cbrt_unary_op<etl::complex<TT>> {
    using T = etl::complex<TT>; 

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

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

    template <typename E>
    static constexpr bool gpu_computable = (is_single_precision_t<T> && impl::egblas::has_scbrt) || (is_double_precision_t<T> && impl::egblas::has_dcbrt)
                                           || (is_complex_single_t<T> && impl::egblas::has_ccbrt) || (is_complex_double_t<T> && impl::egblas::has_zcbrt);

    static constexpr int complexity() {
        return 8;
    }

    static constexpr T apply(const T& x) {
        return etl::cbrt(x);
    }

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

        auto t2 = force_temporary_gpu_dim_only(t1);

        T alpha(1.0);
        impl::egblas::cbrt(etl::size(y), alpha, t1.gpu_memory(), 1, t2.gpu_memory(), 1);

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

        T alpha(1.0);
        impl::egblas::cbrt(etl::size(y), alpha, t1.gpu_memory(), 1, y.gpu_memory(), 1);

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

        return y;
    }

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

} //end of namespace etl