CVPR/LIVE

1#pragma once2 3#include "vector.h"4 5#include <mutex>6#include <condition_variable>7#include <functional>8#include <atomic>9#include <cstdint>10#include <cassert>11#include <algorithm>12// From https://github.com/mmp/pbrt-v3/blob/master/src/core/parallel.h13 14class Barrier {15  public:16    Barrier(int count) : count(count) { assert(count > 0); }17    ~Barrier() { assert(count == 0); }18    void Wait();19 20  private:21    std::mutex mutex;22    std::condition_variable cv;23    int count;24};25 26void parallel_for_host(const std::function<void(int64_t)> &func,27                       int64_t count,28                       int chunkSize = 1);29extern thread_local int ThreadIndex;30void parallel_for_host(31    std::function<void(Vector2i)> func, const Vector2i count);32int num_system_cores();33 34void parallel_init();35void parallel_cleanup();36 37#ifdef __CUDACC__38template <typename T>39__global__ void parallel_for_device_kernel(T functor, int count) {40    auto idx = threadIdx.x + blockIdx.x * blockDim.x;41    if (idx >= count) {42        return;43    }44    functor(idx);45}46template <typename T>47inline void parallel_for_device(T functor,48                                int count,49                                int work_per_thread = 256) {50    if (count <= 0) {51        return;52    }53    auto block_size = work_per_thread;54    auto block_count = idiv_ceil(count, block_size);55    parallel_for_device_kernel<T><<<block_count, block_size>>>(functor, count);56}57#endif58 59template <typename T>60inline void parallel_for(T functor,61                         int count,62                         bool use_gpu,63                         int work_per_thread = -1) {64    if (work_per_thread == -1) {65        work_per_thread = use_gpu ? 64 : 256;66    }67    if (count <= 0) {68        return;69    }70    if (use_gpu) {71#ifdef __CUDACC__72        auto block_size = work_per_thread;73        auto block_count = idiv_ceil(count, block_size);74        parallel_for_device_kernel<T><<<block_count, block_size>>>(functor, count);75#else76        throw std::runtime_error("diffvg not compiled with GPU");77        assert(false);78#endif79    } else {80        auto num_threads = idiv_ceil(count, work_per_thread);81        parallel_for_host([&](int thread_index) {82            auto id_offset = work_per_thread * thread_index;83            auto work_end = std::min(id_offset + work_per_thread, count);84            for (int work_id = id_offset; work_id < work_end; work_id++) {85                auto idx = work_id;86                assert(idx < count);87                functor(idx);88            }89        }, num_threads);90    }91}92