openmp_utils.h

#ifndef STOCHTREE_OPENMP_UTILS_H
#define STOCHTREE_OPENMP_UTILS_H
 
#include <stochtree/common.h>
#include <algorithm>
 
namespace StochTree {
 
#ifdef _OPENMP
 
#include <omp.h>
#define STOCHTREE_HAS_OPENMP 1
 
// OpenMP thread management
inline int get_max_threads() {
    return omp_get_max_threads();
}
 
inline int get_thread_num() {
    return omp_get_thread_num();
}
 
inline int get_num_threads() {
    return omp_get_num_threads();
}
 
inline void set_num_threads(int num_threads) {
    omp_set_num_threads(num_threads);
}
    
#define STOCHTREE_PARALLEL_FOR(num_threads) \
    _Pragma("omp parallel for num_threads(num_threads)")
 
#define STOCHTREE_REDUCTION_ADD(var) \
    _Pragma("omp reduction(+:var)")
 
#define STOCHTREE_CRITICAL \
    _Pragma("omp critical")
 
#else
#define STOCHTREE_HAS_OPENMP 0
 
inline int get_max_threads() {return 1;}
 
inline int get_thread_num() {return 0;}
 
inline int get_num_threads() {return 1;}
 
inline void set_num_threads(int num_threads) {}
    
#define STOCHTREE_PARALLEL_FOR(num_threads)
    
#define STOCHTREE_REDUCTION_ADD(var)
 
#define STOCHTREE_CRITICAL
 
#endif
 
static int GetMaxThreads() {
    return get_max_threads();
}
 
static int GetCurrentThreadNum() {
    return get_thread_num();
}
    
static int GetNumThreads() {
    return get_num_threads();
}
    
static void SetNumThreads(int num_threads) {
    set_num_threads(num_threads);
}
    
static bool IsOpenMPAvailable() {
    return STOCHTREE_HAS_OPENMP;
}
    
static int GetOptimalThreadCount(int workload_size, int min_work_per_thread = 1000) {
    if (!IsOpenMPAvailable()) {
        return 1;
    }
    
    int max_threads = GetMaxThreads();
    int optimal_threads = workload_size / min_work_per_thread;
    
    return std::min(optimal_threads, max_threads);
}
 
// Parallel execution utilities
template<typename Func>
void ParallelFor(int start, int end, int num_threads, Func func) {
    if (num_threads <= 0) {
        num_threads = GetOptimalThreadCount(end - start);
    }
    
    if (num_threads == 1 || !STOCHTREE_HAS_OPENMP) {
        // Sequential execution
        for (int i = start; i < end; ++i) {
            func(i);
        }
    } else {
        // Parallel execution
        STOCHTREE_PARALLEL_FOR(num_threads)
        for (int i = start; i < end; ++i) {
            func(i);
        }
    }
}
 
} // namespace StochTree
 
#endif // STOCHTREE_OPENMP_UTILS_H