mt.hpp

#ifndef MULTI_THREAD_HPP
#define MULTI_THREAD_HPP
#include <chrono>
#include <future>
#include <iostream>
#include <string>
#include <sstream>
#include <iomanip>

#include "def.hpp"
namespace tipl{



class time
{
public:
    // Constructor starts the timer immediately
    time(const char* msg_) : msg(msg_), t1(std::chrono::high_resolution_clock::now()) {}
    time() : t1(std::chrono::high_resolution_clock::now()) {}

    void restart() { t1 = std::chrono::high_resolution_clock::now(); }
    void start() { t1 = std::chrono::high_resolution_clock::now(); }
    void stop() { t2 = std::chrono::high_resolution_clock::now(); }

    template<typename T = std::chrono::milliseconds>
    auto elapsed() {return std::chrono::duration_cast<T>(std::chrono::high_resolution_clock::now() - t1).count();}

    template<typename T = std::chrono::milliseconds>
    auto total() {stop();return std::chrono::duration_cast<T>(t2 - t1).count();}

    std::string to_string() {
        auto now = std::chrono::high_resolution_clock::now();
        auto duration = now - t1;
        auto hours = std::chrono::duration_cast<std::chrono::hours>(duration);
        duration -= hours;
        auto minutes = std::chrono::duration_cast<std::chrono::minutes>(duration);
        duration -= minutes;
        auto seconds = std::chrono::duration_cast<std::chrono::seconds>(duration);
        std::ostringstream oss;
        if (hours.count() > 0) oss << hours.count() << "h";
        if (minutes.count() > 0) oss << minutes.count() << "m";
        oss << seconds.count() << "s";
        return oss.str();
    }

    ~time()
    {
        if (!msg.empty())
            std::cout << msg << elapsed<>() << std::endl;
    }

private:
    std::string msg;
    std::chrono::high_resolution_clock::time_point t1, t2;
};

class estimate_time{
    std::string name;
    size_t n = 0;
    double time_total = 0.0;
    std::chrono::high_resolution_clock::time_point s;
public:
    estimate_time(const char* name_):name(name_){}
    ~estimate_time()
    {
        if(n)
            std::cout << name << time_total/double(n) << " microseconds" << std::endl;
    }
    void start(void)
    {
        s = std::chrono::high_resolution_clock::now();
    }
    void stop(void)
    {
        auto stop = std::chrono::high_resolution_clock::now();
        time_total += std::chrono::duration_cast<std::chrono::microseconds>(stop-s).count();
        ++n;
    }
};


inline auto main_thread_id = std::this_thread::get_id();
inline bool is_main_thread(void)
{
    return main_thread_id == std::this_thread::get_id();
}

inline int max_thread_count = std::thread::hardware_concurrency();


enum par_for_type{
    sequential = 0,
    sequential_with_id = 1,
    ranged = 2,
    ranged_with_id = 3,
    dynamic = 4,
    dynamic_with_id = 5
};

inline std::atomic<bool> par_for_running = false;
struct par_for_guard
{
    bool is_root;
    par_for_guard(void)
    {
        is_root = !par_for_running.exchange(true);
    }
    ~par_for_guard(void)
    {
        if(is_root)
            par_for_running = false;
    }
};
template <par_for_type type = dynamic,typename T,
          typename Func,typename std::enable_if<
              std::is_integral<T>::value ||
              std::is_class<T>::value ||
              std::is_pointer<T>::value,bool>::type = true>
__HOST__ void par_for(T from,T to,Func&& f,int thread_count)
{
    if (from == to)
            return;

    size_t n = to - from;
    thread_count = std::min<int>(thread_count, (int)n);
    par_for_guard guard;
    if(thread_count >= max_thread_count)
        thread_count = (guard.is_root && n > 1) ?  thread_count : 1;


#ifdef __CUDACC__
    int cur_device = 0;
    bool set_cuda_dev = false;
    if constexpr(use_cuda)
    {
        set_cuda_dev = (thread_count > 1 && cudaGetDevice(&cur_device) == cudaSuccess);
    }
#endif

    std::vector<std::thread> workers;

    if constexpr(type == dynamic || type == dynamic_with_id)
    {
        std::atomic<size_t> next_idx{0};
        auto run = [=, &f, &next_idx](size_t id)
        {
#ifdef __CUDACC__
            if constexpr(use_cuda)
            {
                if(id && set_cuda_dev)
                    cudaSetDevice(cur_device);
            }
#endif
            for (size_t i = next_idx++; i < n; i = next_idx++)
            {
                if constexpr (type == dynamic_with_id)
                    f(from + i, id);
                else
                    f(from + i);
            }
        };

        for (int i = 1; i < thread_count; ++i)
            workers.emplace_back(run, i);

        run(0);
    }
    else
    {
        auto run = [=, &f](size_t id)
        {
#ifdef __CUDACC__
            if constexpr(use_cuda)
            {
                if(id && set_cuda_dev)
                    cudaSetDevice(cur_device);
            }
#endif
            size_t block = n / thread_count, rem = n % thread_count;
            T beg = from + (id * block + std::min<size_t>(id, rem));
            T end = beg + block + (id < rem);

            if constexpr(type == sequential || type == sequential_with_id)
            {
                for(; beg != end; ++beg)
                {
                    if constexpr(type == sequential_with_id)
                        f(beg, id);
                    else
                        f(beg);
                }
            }
            else if constexpr(type == ranged || type == ranged_with_id)
            {
                if constexpr(type == ranged_with_id)
                    f(beg, end, id);
                else
                    f(beg, end);
            }
        };

        for (int i = 1; i < thread_count; ++i)
            workers.emplace_back(run, i);

        run(0);
    }

    for (auto& t : workers)
        t.join();

}

template <par_for_type type = dynamic, typename T, typename Func,
          typename std::enable_if_t<std::is_integral_v<T> || std::is_class_v<T> || std::is_pointer_v<T>, int> = 0>
inline void par_for(T from, T to, Func&& f) {
    par_for<type>(from, to, std::forward<Func>(f), par_for_running ? 1 : max_thread_count);
}

template <par_for_type type = dynamic, typename T, typename Func, typename std::enable_if_t<std::is_integral_v<T>, int> = 0>
inline void par_for(T size, Func&& f, int tc = max_thread_count) {
    par_for<type>(T(0), size, std::forward<Func>(f), tc);
}



template <typename T>
inline double estimate_run_time(T&& fun)
{
    auto start = std::chrono::steady_clock::now();
    std::forward<T>(fun)();
    auto end = std::chrono::steady_clock::now();

    return std::chrono::duration<double, std::micro>(end - start).count();
}

template<typename T>
void aggregate_results(std::vector<std::vector<T> >&& results,std::vector<T>& all_result_)
{
    std::vector<size_t> insert_pos;
    insert_pos.push_back(0);
    for(size_t i = 0,sz = results.size();i < sz;++i)
        insert_pos.push_back(insert_pos.back() + results[i].size());

    std::vector<T> all_result(insert_pos.back());
    tipl::par_for(results.size(),[&](unsigned int index)
    {
        std::move(results[index].begin(),results[index].end(),all_result.begin()+int64_t(insert_pos[index]));
    });
    all_result.swap(all_result_);
}


class thread{
    std::unique_ptr<std::thread> th;
public:
    bool running = false;
    bool terminated = false;
    #ifdef __CUDACC__
    int cur_device = 0;
    #endif
public:
    thread(void){}
    ~thread(void)
    {
        #ifdef __CUDACC__
        if constexpr(use_cuda)
            cudaDeviceSynchronize();
        #endif
        clear();
    }
    void clear(void)
    {
        if(th)
        {
            terminated = true;
            th->join();
            th.reset();
        }
        terminated = false;
        running = false;
    }
    template<typename lambda_type>
    void run(lambda_type&& fun)
    {
        clear();
        #ifdef __CUDACC__
        if constexpr(use_cuda)
        {
            if(cudaGetDevice(&cur_device) != cudaSuccess)
                throw std::runtime_error(cudaGetErrorName(cudaGetLastError()));
        }
        #endif
        th = std::make_unique<std::thread>([this,fun = std::forward<lambda_type>(fun)]()
        {
            running = true;
            #ifdef __CUDACC__
            if constexpr(use_cuda)
            {
                if(cudaSetDevice(cur_device) != cudaSuccess)
                    throw std::runtime_error(cudaGetErrorName(cudaGetLastError()));
            }
            #endif
            try{
            fun();
            }
            catch(...){
                running = false;
                throw;
            }
            running = false;
        });
    }
    void join(void)
    {
        th->join();
    }
};


}
#endif // MULTI_THREAD_HPP