tor-browser

resizable_parallel_runner.cc (6014B)

---

// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include <jxl/jxl_threads_export.h>
#include <jxl/memory_manager.h>
#include <jxl/parallel_runner.h>
#include <jxl/resizable_parallel_runner.h>

#include <algorithm>
#include <atomic>
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <mutex>
#include <thread>
#include <vector>

namespace jpegxl {
namespace {

// A thread pool that allows changing the number of threads it runs. It also
// runs tasks on the calling thread, which can work better on schedulers for
// heterogeneous architectures.
struct ResizeableParallelRunner {
 void SetNumThreads(size_t num) {
   if (num > 0) {
     num -= 1;
   }
   {
     std::unique_lock<std::mutex> l(state_mutex_);
     num_desired_workers_ = num;
     workers_can_proceed_.notify_all();
   }
   if (workers_.size() < num) {
     for (size_t i = workers_.size(); i < num; i++) {
       workers_.emplace_back([this, i]() { WorkerBody(i); });
     }
   }
   if (workers_.size() > num) {
     for (size_t i = num; i < workers_.size(); i++) {
       workers_[i].join();
     }
     workers_.resize(num);
   }
 }

 ~ResizeableParallelRunner() { SetNumThreads(0); }

 JxlParallelRetCode Run(void* jxl_opaque, JxlParallelRunInit init,
                        JxlParallelRunFunction func, uint32_t start,
                        uint32_t end) {
   if (start + 1 == end) {
     JxlParallelRetCode ret = init(jxl_opaque, 1);
     if (ret != 0) return ret;

     func(jxl_opaque, start, 0);
     return ret;
   }

   size_t num_workers = std::min<size_t>(workers_.size() + 1, end - start);
   JxlParallelRetCode ret = init(jxl_opaque, num_workers);
   if (ret != 0) {
     return ret;
   }

   {
     std::unique_lock<std::mutex> l(state_mutex_);
     // Avoid waking up more workers than needed.
     max_running_workers_ = end - start - 1;
     next_task_ = start;
     end_task_ = end;
     func_ = func;
     jxl_opaque_ = jxl_opaque;
     work_available_ = true;
     num_running_workers_++;
     workers_can_proceed_.notify_all();
   }

   DequeueTasks(0);

   while (true) {
     std::unique_lock<std::mutex> l(state_mutex_);
     if (num_running_workers_ == 0) break;
     work_done_.wait(l);
   }

   return ret;
 }

private:
 void WorkerBody(size_t worker_id) {
   while (true) {
     {
       std::unique_lock<std::mutex> l(state_mutex_);
       // Worker pool was reduced, resize down.
       if (worker_id >= num_desired_workers_) {
         return;
       }
       // Nothing to do this time.
       if (!work_available_ || worker_id >= max_running_workers_) {
         workers_can_proceed_.wait(l);
         continue;
       }
       num_running_workers_++;
     }
     DequeueTasks(worker_id + 1);
   }
 }

 void DequeueTasks(size_t thread_id) {
   while (true) {
     uint32_t task = next_task_++;
     if (task >= end_task_) {
       std::unique_lock<std::mutex> l(state_mutex_);
       num_running_workers_--;
       work_available_ = false;
       if (num_running_workers_ == 0) {
         work_done_.notify_all();
       }
       break;
     }
     func_(jxl_opaque_, task, thread_id);
   }
 }

 // Checks when the worker has something to do, which can be one of:
 // - quitting (when worker_id >= num_desired_workers_)
 // - having work available for them (work_available_ is true and worker_id >=
 // max_running_workers_)
 std::condition_variable workers_can_proceed_;

 // Workers are done, and the main thread can proceed (num_running_workers_ ==
 // 0)
 std::condition_variable work_done_;

 std::vector<std::thread> workers_;

 // Protects all the remaining variables, except for func_, jxl_opaque_ and
 // end_task_ (for which only the write by the main thread is protected, and
 // subsequent uses by workers happen-after it) and next_task_ (which is
 // atomic).
 std::mutex state_mutex_;

 // Range of tasks still need to be done.
 std::atomic<uint32_t> next_task_;
 uint32_t end_task_;

 // Function to run and its argument.
 JxlParallelRunFunction func_;
 void* jxl_opaque_;  // not owned

 // Variables that control the workers:
 // - work_available_ is set to true after a call to Run() and to false at the
 // end of it.
 // - num_desired_workers_ represents the number of workers that should be
 // present.
 // - max_running_workers_ represents the number of workers that should be
 // executing tasks.
 // - num_running_workers_ represents the number of workers that are executing
 // tasks.
 size_t num_desired_workers_ = 0;
 size_t max_running_workers_ = 0;
 size_t num_running_workers_ = 0;
 bool work_available_ = false;
};
}  // namespace
}  // namespace jpegxl

extern "C" {
JXL_THREADS_EXPORT JxlParallelRetCode JxlResizableParallelRunner(
   void* runner_opaque, void* jpegxl_opaque, JxlParallelRunInit init,
   JxlParallelRunFunction func, uint32_t start_range, uint32_t end_range) {
 return static_cast<jpegxl::ResizeableParallelRunner*>(runner_opaque)
     ->Run(jpegxl_opaque, init, func, start_range, end_range);
}

JXL_THREADS_EXPORT void* JxlResizableParallelRunnerCreate(
   const JxlMemoryManager* memory_manager) {
 return new jpegxl::ResizeableParallelRunner();
}

JXL_THREADS_EXPORT void JxlResizableParallelRunnerSetThreads(
   void* runner_opaque, size_t num_threads) {
 static_cast<jpegxl::ResizeableParallelRunner*>(runner_opaque)
     ->SetNumThreads(num_threads);
}

JXL_THREADS_EXPORT void JxlResizableParallelRunnerDestroy(void* runner_opaque) {
 delete static_cast<jpegxl::ResizeableParallelRunner*>(runner_opaque);
}

JXL_THREADS_EXPORT uint32_t
JxlResizableParallelRunnerSuggestThreads(uint64_t xsize, uint64_t ysize) {
 // ~one thread per group.
 return std::min<uint64_t>(std::thread::hardware_concurrency(),
                           xsize * ysize / (256 * 256));
}
}