context.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 {

#ifdef ETL_MANUAL_SELECT

template <typename T>
struct forced_impl {
    T impl;              
    bool forced = false; 
};
#endif

struct context {
    bool serial   = false; 
    bool parallel = false; 
    bool cpu      = false; 

#ifdef ETL_MANUAL_SELECT
    forced_impl<sum_impl> sum_selector;               
    forced_impl<pool_impl> pool_selector;             
    forced_impl<transpose_impl> transpose_selector;   
    forced_impl<dot_impl> dot_selector;               
    forced_impl<conv_impl> conv_selector;             
    forced_impl<conv_multi_impl> conv_multi_selector; 
    forced_impl<conv4_impl> conv4_selector;           
    forced_impl<gemm_impl> gemm_selector;             
    forced_impl<outer_impl> outer_selector;           
    forced_impl<bias_add_impl> bias_add_selector;     
    forced_impl<fft_impl> fft_selector;               
#endif
};

inline context& local_context() {
    static thread_local context local_context;
    return local_context;
}

inline bool is_something_forced() {
#ifdef ETL_MANUAL_SELECT
    auto& c = local_context();
    return c.sum_selector.forced || c.pool_selector.forced || c.transpose_selector.forced || c.dot_selector.forced || c.conv_selector.forced
           || c.conv_multi_selector.forced || c.conv4_selector.forced || c.gemm_selector.forced || c.outer_selector.forced || c.bias_add_selector.forced
           || c.fft_selector.forced;
#else
    return false;
#endif
}

namespace detail {

#ifdef ETL_MANUAL_SELECT

template <typename T>
forced_impl<T>& get_forced_impl();

template <>
inline forced_impl<sum_impl>& get_forced_impl() {
    return local_context().sum_selector;
}

template <>
inline forced_impl<pool_impl>& get_forced_impl() {
    return local_context().pool_selector;
}

template <>
inline forced_impl<transpose_impl>& get_forced_impl() {
    return local_context().transpose_selector;
}

template <>
inline forced_impl<dot_impl>& get_forced_impl() {
    return local_context().dot_selector;
}

template <>
inline forced_impl<conv_impl>& get_forced_impl() {
    return local_context().conv_selector;
}

template <>
inline forced_impl<conv_multi_impl>& get_forced_impl() {
    return local_context().conv_multi_selector;
}

template <>
inline forced_impl<conv4_impl>& get_forced_impl() {
    return local_context().conv4_selector;
}

template <>
inline forced_impl<gemm_impl>& get_forced_impl() {
    return local_context().gemm_selector;
}

template <>
inline forced_impl<outer_impl>& get_forced_impl() {
    return local_context().outer_selector;
}

template <>
inline forced_impl<bias_add_impl>& get_forced_impl() {
    return local_context().bias_add_selector;
}

template <>
inline forced_impl<fft_impl>& get_forced_impl() {
    return local_context().fft_selector;
}

#endif

struct serial_context {
    bool old_serial; 

    serial_context() {
        old_serial                  = etl::local_context().serial;
        etl::local_context().serial = true;
    }

    ~serial_context() {
        etl::local_context().serial = old_serial;
    }

    operator bool() {
        return true;
    }
};

struct parallel_context {
    bool old_parallel; 

    parallel_context() {
        old_parallel                  = etl::local_context().parallel;
        etl::local_context().parallel = true;
    }

    ~parallel_context() {
        etl::local_context().parallel = old_parallel;
    }

    operator bool() {
        return true;
    }
};

struct cpu_context {
    bool old_cpu; 

    cpu_context() {
        old_cpu                  = etl::local_context().cpu;
        etl::local_context().cpu = true;
    }

    ~cpu_context() {
        etl::local_context().cpu = old_cpu;
    }

    operator bool() {
        return true;
    }
};

#ifdef ETL_MANUAL_SELECT

template <typename Selector, Selector V>
struct selected_context {
    forced_impl<Selector> old_selector; 

    selected_context() {
        decltype(auto) selector = get_forced_impl<Selector>();

        old_selector = selector;

        selector.impl   = V;
        selector.forced = true;
    }

    ~selected_context() {
        get_forced_impl<Selector>() = old_selector;
    }

    operator bool() {
        return true;
    }
};

#endif

} //end of namespace detail

#define SERIAL_SECTION if (auto etl_serial_context__ = etl::detail::serial_context())

#define PARALLEL_SECTION if (auto etl_parallel_context__ = etl::detail::parallel_context())

#define CPU_SECTION if (auto etl_cpu_context__ = etl::detail::cpu_context())

#ifdef ETL_MANUAL_SELECT

#define SELECTED_SECTION(v) if (auto etl_selected_context__ = etl::detail::selected_context<decltype(v), v>())

#endif

} //end of namespace etl