thread_pool.hpp

Go to the documentation of this file.

//This file is part of Bertini 2.
//
//bertini2/parallel/thread_pool.hpp is free software: you can redistribute it and/or modify
//it under the terms of the GNU General Public License as published by
//the Free Software Foundation, either version 3 of the License, or
//(at your option) any later version.
//
//bertini2/parallel/thread_pool.hpp is distributed in the hope that it will be useful,
//but WITHOUT ANY WARRANTY; without even the implied warranty of
//MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//GNU General Public License for more details.
//
//You should have received a copy of the GNU General Public License
//along with bertini2/parallel/thread_pool.hpp.  If not, see <http://www.gnu.org/licenses/>.
//
// Copyright(C) Bertini2 Development Team

 
#pragma once
 
// NOTE: this header is intentionally NOT gated behind BERTINI2_HAVE_MPI.  The thread pool is
// pure <thread>/<mutex>/<condition_variable> with no MPI dependency; it backs both the
// MPI+threads hybrid worker AND the MPI-less shared-memory threaded solve.  See ADR / the
// "MPI-less threading" plan: threading must be available in builds (and wheels) compiled
// without any MPI installation.
 
#if defined(__linux__)
#include <sched.h>   // sched_getaffinity, CPU_COUNT (AvailableCpuCount)
#endif
 
#include <condition_variable>
#include <cstdlib>   // std::getenv, std::atoi
#include <deque>
#include <functional>
#include <mutex>
#include <optional>
#include <thread>
#include <type_traits>
#include <variant>
#include <vector>
 
namespace bertini {
namespace parallel {
 

inline unsigned AvailableCpuCount()
{
#if defined(__linux__)
    cpu_set_t mask;
    CPU_ZERO(&mask);
    if (sched_getaffinity(0, sizeof(mask), &mask) == 0)
    {
        const int n = CPU_COUNT(&mask);
        if (n >= 1)
            return static_cast<unsigned>(n);
    }
#endif
    const unsigned hw = std::thread::hardware_concurrency();
    return hw >= 1 ? hw : 1u;
}
 

// b2 reads its own variable, never OMP_NUM_THREADS: numpy's OpenBLAS takes its thread count from
// that one, so a shared name would tie the two libraries' thread counts together (ADR-0070).
inline unsigned EffectiveThreadCount(unsigned configured = 0)
{
    if (const char* env = std::getenv("BERTINI_NUM_THREADS"))
    {
        int n = std::atoi(env);
        if (n >= 1)
            return static_cast<unsigned>(n);
    }
 
    if (configured >= 1)
        return configured;
 
    return AvailableCpuCount();
}
 

template<typename T>
class ThreadSafeQueue
{
    std::deque<T>           queue_;
    std::mutex              mutex_;
    std::condition_variable cv_;
 
public:
    void push(T item)
    {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            queue_.push_back(std::move(item));
        }
        cv_.notify_one();
    }

    T pop()
    {
        std::unique_lock<std::mutex> lock(mutex_);
        cv_.wait(lock, [this]{ return !queue_.empty(); });
        T item = std::move(queue_.front());
        queue_.pop_front();
        return item;
    }

    std::optional<T> try_pop()
    {
        std::lock_guard<std::mutex> lock(mutex_);
        if (queue_.empty())
            return std::nullopt;
        T item = std::move(queue_.front());
        queue_.pop_front();
        return item;
    }
};
 
 
struct PoolShutdownSentinel {};
 

template<typename TaskT, typename ResultT, typename StateFactory, typename TrackFn>
class WorkerThreadPool
{
    using StateT   = std::invoke_result_t<StateFactory>;
    using WorkItem = std::variant<TaskT, PoolShutdownSentinel>;
 
    ThreadSafeQueue<WorkItem>  task_queue_;
    ThreadSafeQueue<ResultT>   result_queue_;
    std::vector<std::thread>   threads_;
    int                        n_threads_;
 
public:
    WorkerThreadPool(int n_threads, StateFactory state_factory, TrackFn track_fn)
        : n_threads_(n_threads)
    {
        threads_.reserve(static_cast<size_t>(n_threads));
        for (int i = 0; i < n_threads; ++i)
        {
            threads_.emplace_back([this, state_factory, track_fn]() mutable
            {
                // Build thread-local state once (System copy + Tracker copy).
                StateT state = state_factory();
 
                while (true)
                {
                    WorkItem item = task_queue_.pop();
 
                    if (std::holds_alternative<PoolShutdownSentinel>(item))
                        break;
 
                    TaskT const& task = std::get<TaskT>(item);
                    ResultT result = track_fn(state, task);
                    result_queue_.push(std::move(result));
                }
            });
        }
    }

    void submit(TaskT task)
    {
        task_queue_.push(WorkItem{std::move(task)});
    }

    ResultT collect()
    {
        return result_queue_.pop();
    }

    std::optional<ResultT> try_collect()
    {
        return result_queue_.try_pop();
    }

    void shutdown()
    {
        for (int i = 0; i < n_threads_; ++i)
            task_queue_.push(WorkItem{PoolShutdownSentinel{}});
        for (auto& t : threads_)
            t.join();
        threads_.clear();
    }
};
 
 
 
} // namespace parallel
} // namespace bertini