Method ParallelForImpl

oneflow/core/thread/thread_runtime.h:73–93

class OfRuntime final : public RuntimeBase {
 private:
  void ParallelForImpl(int64_t begin, int64_t end, const CallableT& func, size_t num_threads,
                       size_t grain_size) override {
    if (unlikely(pthread_fork::IsForkedSubProcess()) || Singleton<ThreadPool>::Get() == nullptr) {
      return SeqFor(begin, end, func);
    }
    const size_t num_elements = end - begin;
    num_threads = std::min(num_elements, num_threads);
    BalancedSplitter bs(num_elements, num_threads);
    BlockingCounter bc(num_threads);

    FOR_RANGE(size_t, range_id, 0, num_threads) {
      Singleton<ThreadPool>::Get()->AddWork([&bc, &bs, range_id, func] {
        const size_t begin_ = bs.At(range_id).begin();
        const size_t end_ = bs.At(range_id).end();
        SeqFor(begin_, end_, func);
        bc.Decrease();
      });
    }
    // buzy loop wait.
    bc.WaitForeverUntilCntEqualZero();
  }
};

#if WITH_TBB