92 lines
2.5 KiB
C++
92 lines
2.5 KiB
C++
#pragma once
|
|
|
|
#include "vector.h"
|
|
|
|
#include <mutex>
|
|
#include <condition_variable>
|
|
#include <functional>
|
|
#include <atomic>
|
|
#include <cstdint>
|
|
#include <cassert>
|
|
#include <algorithm>
|
|
// From https://github.com/mmp/pbrt-v3/blob/master/src/core/parallel.h
|
|
|
|
class Barrier {
|
|
public:
|
|
Barrier(int count) : count(count) { assert(count > 0); }
|
|
~Barrier() { assert(count == 0); }
|
|
void Wait();
|
|
|
|
private:
|
|
std::mutex mutex;
|
|
std::condition_variable cv;
|
|
int count;
|
|
};
|
|
|
|
void parallel_for_host(const std::function<void(int64_t)> &func,
|
|
int64_t count,
|
|
int chunkSize = 1);
|
|
extern thread_local int ThreadIndex;
|
|
void parallel_for_host(
|
|
std::function<void(Vector2i)> func, const Vector2i count);
|
|
int num_system_cores();
|
|
|
|
void parallel_init();
|
|
void parallel_cleanup();
|
|
|
|
#ifdef __CUDACC__
|
|
template <typename T>
|
|
__global__ void parallel_for_device_kernel(T functor, int count) {
|
|
auto idx = threadIdx.x + blockIdx.x * blockDim.x;
|
|
if (idx >= count) {
|
|
return;
|
|
}
|
|
functor(idx);
|
|
}
|
|
template <typename T>
|
|
inline void parallel_for_device(T functor,
|
|
int count,
|
|
int work_per_thread = 256) {
|
|
if (count <= 0) {
|
|
return;
|
|
}
|
|
auto block_size = work_per_thread;
|
|
auto block_count = idiv_ceil(count, block_size);
|
|
parallel_for_device_kernel<T><<<block_count, block_size>>>(functor, count);
|
|
}
|
|
#endif
|
|
|
|
template <typename T>
|
|
inline void parallel_for(T functor,
|
|
int count,
|
|
bool use_gpu,
|
|
int work_per_thread = -1) {
|
|
if (work_per_thread == -1) {
|
|
work_per_thread = use_gpu ? 64 : 256;
|
|
}
|
|
if (count <= 0) {
|
|
return;
|
|
}
|
|
if (use_gpu) {
|
|
#ifdef __CUDACC__
|
|
auto block_size = work_per_thread;
|
|
auto block_count = idiv_ceil(count, block_size);
|
|
parallel_for_device_kernel<T><<<block_count, block_size>>>(functor, count);
|
|
#else
|
|
throw std::runtime_error("diffvg not compiled with GPU");
|
|
assert(false);
|
|
#endif
|
|
} else {
|
|
auto num_threads = idiv_ceil(count, work_per_thread);
|
|
parallel_for_host([&](int thread_index) {
|
|
auto id_offset = work_per_thread * thread_index;
|
|
auto work_end = std::min(id_offset + work_per_thread, count);
|
|
for (int work_id = id_offset; work_id < work_end; work_id++) {
|
|
auto idx = work_id;
|
|
assert(idx < count);
|
|
functor(idx);
|
|
}
|
|
}, num_threads);
|
|
}
|
|
}
|