helpers.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 E>
constexpr size_t dimensions([[maybe_unused]] const E& expr) noexcept {
    return etl_traits<E>::dimensions();
}

template <typename E>
constexpr size_t dimensions() noexcept {
    return decay_traits<E>::dimensions();
}

template <typename E>
constexpr int complexity([[maybe_unused]] const E& expr) noexcept {
    return etl_traits<E>::complexity();
}

template <typename E>
constexpr int complexity() noexcept {
    return decay_traits<E>::complexity();
}

template <dyn_expr E>
size_t rows(const E& expr) {
    return etl_traits<E>::dim(expr, 0);
}

template <fast_expr E>
constexpr size_t rows(const E& expr) noexcept {
    return (void)expr, etl_traits<E>::template dim<0>();
}

template <dyn_matrix_c E>
size_t columns(const E& expr) {
    return etl_traits<E>::dim(expr, 1);
}

template <fast_matrix_c E>
constexpr size_t columns(const E& expr) noexcept {
    return (void)expr, etl_traits<E>::template dim<1>();
}

template <dyn_expr E>
size_t size(const E& expr) {
    return etl_traits<E>::size(expr);
}

template <fast_expr E>
constexpr size_t size(const E& expr) noexcept {
    return (void)expr, etl_traits<E>::size();
}

template <dyn_matrix_c E>
size_t subsize(const E& expr) {
    return etl_traits<E>::size(expr) / etl_traits<E>::dim(expr, 0);
}

template <fast_matrix_c E>
constexpr size_t subsize(const E& expr) noexcept {
    return (void)expr, etl_traits<E>::size() / etl_traits<E>::template dim<0>();
}

template <size_t D, dyn_expr E>
size_t dim(const E& e) noexcept {
    return etl_traits<E>::dim(e, D);
}

template <etl_expr E>
size_t dim(const E& e, size_t d) noexcept {
    return etl_traits<E>::dim(e, d);
}

template <size_t D, fast_expr E>
constexpr size_t dim(const E& e) noexcept {
    return (void)e, etl_traits<E>::template dim<D>();
}

template <size_t D, fast_expr E>
constexpr size_t dim() noexcept {
    return decay_traits<E>::template dim<D>();
}

template <typename E>
struct safe_dimensions_impl;

template <typename E>
requires(generator<E>)
struct safe_dimensions_impl<E>  : std::integral_constant<size_t, std::numeric_limits<size_t>::max()> {};

template <typename E>
requires(!generator<E>)
struct safe_dimensions_impl<E>  : std::integral_constant<size_t, etl_traits<E>::dimensions()> {};

template <typename E, typename Enable = void>
constexpr size_t safe_dimensions = safe_dimensions_impl<E>::value;

template <typename E>
constexpr std::pair<size_t, size_t> index_to_2d(E&& sub, size_t i) {
    return decay_traits<E>::storage_order == order::RowMajor ? std::make_pair(i / dim<1>(sub), i % dim<1>(sub))
                                                             : std::make_pair(i % dim<0>(sub), i / dim<0>(sub));
}

template <etl_2d E>
size_t row_stride(E&& expr) {
    return decay_traits<E>::storage_order == order::RowMajor ? etl::dim<1>(expr) : 1;
}

template <etl_2d E>
size_t col_stride(E&& expr) {
    return decay_traits<E>::storage_order == order::RowMajor ? 1 : etl::dim<0>(expr);
}

template <etl_2d E>
size_t minor_stride(E&& expr) {
    return decay_traits<E>::storage_order == order::RowMajor ? etl::dim<0>(expr) : etl::dim<1>(expr);
}

template <etl_2d E>
size_t major_stride(E&& expr) {
    return decay_traits<E>::storage_order == order::RowMajor ? etl::dim<1>(expr) : etl::dim<0>(expr);
}

template <typename P1, typename P2>
bool memory_alias(const P1* a_begin, const P1* a_end, const P2* b_begin, const P2* b_end) {
    cpp_assert(a_begin <= a_end, "memory_alias works on ordered ranges");
    cpp_assert(b_begin <= b_end, "memory_alias works on ordered ranges");

    return reinterpret_cast<uintptr_t>(a_begin) < reinterpret_cast<uintptr_t>(b_end)
           && reinterpret_cast<uintptr_t>(a_end) > reinterpret_cast<uintptr_t>(b_begin);
}

template <typename E>
void safe_ensure_cpu_up_to_date(E&& expr) {
    expr.ensure_cpu_up_to_date();
}

template <typename E>
bool safe_is_cpu_up_to_date(E&& expr) {
    if constexpr (is_dma<E>) {
        return expr.is_cpu_up_to_date();
    } else {
        return true;
    }
}

template <typename E>
bool safe_is_gpu_up_to_date(E&& expr) {
    if constexpr (is_dma<E>) {
        return expr.is_gpu_up_to_date();
    } else {
        return false;
    }
}

template <typename E>
decltype(auto) smart_forward(E& expr) {
    if constexpr (is_temporary_expr<E>) {
        return force_temporary(expr);
    } else {
        return make_temporary(expr);
    }
}

template <typename E>
decltype(auto) smart_forward_gpu(E& expr) {
    if constexpr (is_temporary_expr<E>) {
        if constexpr (E::gpu_computable) {
            return force_temporary_gpu(expr);
        } else {
            return force_temporary(expr);
        }
    } else {
        return make_temporary(expr);
    }
}

// Unary smart_gpu_compute

template <typename E, typename Y>
decltype(auto) smart_gpu_compute_hint(E& expr, Y& y) {
    if constexpr (is_temporary_expr<E>) {
        if constexpr (E::gpu_computable) {
            return force_temporary_gpu(expr);
        } else {
            auto t = force_temporary(expr);
            t.ensure_gpu_up_to_date();
            return t;
        }
    } else {
        return expr.gpu_compute_hint(y);
    }
}

// Binary smart_gpu_compute

template <typename X, typename Y>
decltype(auto) smart_gpu_compute(X& x, Y& y) {
    if constexpr (is_temporary_expr<X>) {
        if constexpr (X::gpu_computable) {
            return y = x;
        } else {
            auto t = force_temporary(x);
            t.ensure_gpu_up_to_date();
            y = t;
            return y;
        }
    } else {
        if constexpr (is_dma<X>) {
            x.ensure_gpu_up_to_date();
            y.ensure_gpu_allocated();
            y.gpu_copy_from(x.gpu_memory());
            return y;
        } else {
            return x.gpu_compute(y);
        }
    }
}

// Select version

template <typename X, typename Y>
decltype(auto) select_smart_gpu_compute(X& x, Y& y) {
    if constexpr (should_gpu_compute_direct<X>) {
        return smart_gpu_compute_hint(x, y);
    } else {
        return smart_gpu_compute(x, y);
    }
}

template <typename T>
constexpr size_t gpu_inc = is_scalar<T> ? 0 : 1;

} //end of namespace etl