Androidonnxfork's picture
Upload folder using huggingface_hub
8b7c501
/* Standard C headers */
#include <assert.h>
#include <limits.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
/* Configuration header */
#include "threadpool-common.h"
/* POSIX headers */
#include <pthread.h>
#include <unistd.h>
/* Futex-specific headers */
#if PTHREADPOOL_USE_FUTEX
#if defined(__linux__)
#include <sys/syscall.h>
#include <linux/futex.h>
/* Old Android NDKs do not define SYS_futex and FUTEX_PRIVATE_FLAG */
#ifndef SYS_futex
#define SYS_futex __NR_futex
#endif
#ifndef FUTEX_PRIVATE_FLAG
#define FUTEX_PRIVATE_FLAG 128
#endif
#elif defined(__EMSCRIPTEN__)
/* math.h for INFINITY constant */
#include <math.h>
#include <emscripten/threading.h>
#else
#error "Platform-specific implementation of futex_wait and futex_wake_all required"
#endif
#endif
/* Windows-specific headers */
#ifdef _WIN32
#include <sysinfoapi.h>
#endif
/* Dependencies */
#if PTHREADPOOL_USE_CPUINFO
#include <cpuinfo.h>
#endif
/* Public library header */
#include <pthreadpool.h>
/* Internal library headers */
#include "threadpool-atomics.h"
#include "threadpool-object.h"
#include "threadpool-utils.h"
#if PTHREADPOOL_USE_FUTEX
#if defined(__linux__)
static int futex_wait(pthreadpool_atomic_uint32_t* address, uint32_t value) {
return syscall(SYS_futex, address, FUTEX_WAIT | FUTEX_PRIVATE_FLAG, value, NULL);
}
static int futex_wake_all(pthreadpool_atomic_uint32_t* address) {
return syscall(SYS_futex, address, FUTEX_WAKE | FUTEX_PRIVATE_FLAG, INT_MAX);
}
#elif defined(__EMSCRIPTEN__)
static int futex_wait(pthreadpool_atomic_uint32_t* address, uint32_t value) {
return emscripten_futex_wait((volatile void*) address, value, INFINITY);
}
static int futex_wake_all(pthreadpool_atomic_uint32_t* address) {
return emscripten_futex_wake((volatile void*) address, INT_MAX);
}
#else
#error "Platform-specific implementation of futex_wait and futex_wake_all required"
#endif
#endif
static void checkin_worker_thread(struct pthreadpool* threadpool) {
#if PTHREADPOOL_USE_FUTEX
if (pthreadpool_decrement_fetch_acquire_release_size_t(&threadpool->active_threads) == 0) {
pthreadpool_store_release_uint32_t(&threadpool->has_active_threads, 0);
futex_wake_all(&threadpool->has_active_threads);
}
#else
pthread_mutex_lock(&threadpool->completion_mutex);
if (pthreadpool_decrement_fetch_release_size_t(&threadpool->active_threads) == 0) {
pthread_cond_signal(&threadpool->completion_condvar);
}
pthread_mutex_unlock(&threadpool->completion_mutex);
#endif
}
static void wait_worker_threads(struct pthreadpool* threadpool) {
/* Initial check */
#if PTHREADPOOL_USE_FUTEX
uint32_t has_active_threads = pthreadpool_load_acquire_uint32_t(&threadpool->has_active_threads);
if (has_active_threads == 0) {
return;
}
#else
size_t active_threads = pthreadpool_load_acquire_size_t(&threadpool->active_threads);
if (active_threads == 0) {
return;
}
#endif
/* Spin-wait */
for (uint32_t i = PTHREADPOOL_SPIN_WAIT_ITERATIONS; i != 0; i--) {
pthreadpool_yield();
#if PTHREADPOOL_USE_FUTEX
has_active_threads = pthreadpool_load_acquire_uint32_t(&threadpool->has_active_threads);
if (has_active_threads == 0) {
return;
}
#else
active_threads = pthreadpool_load_acquire_size_t(&threadpool->active_threads);
if (active_threads == 0) {
return;
}
#endif
}
/* Fall-back to mutex/futex wait */
#if PTHREADPOOL_USE_FUTEX
while ((has_active_threads = pthreadpool_load_acquire_uint32_t(&threadpool->has_active_threads)) != 0) {
futex_wait(&threadpool->has_active_threads, 1);
}
#else
pthread_mutex_lock(&threadpool->completion_mutex);
while (pthreadpool_load_acquire_size_t(&threadpool->active_threads) != 0) {
pthread_cond_wait(&threadpool->completion_condvar, &threadpool->completion_mutex);
};
pthread_mutex_unlock(&threadpool->completion_mutex);
#endif
}
static uint32_t wait_for_new_command(
struct pthreadpool* threadpool,
uint32_t last_command,
uint32_t last_flags)
{
uint32_t command = pthreadpool_load_acquire_uint32_t(&threadpool->command);
if (command != last_command) {
return command;
}
if ((last_flags & PTHREADPOOL_FLAG_YIELD_WORKERS) == 0) {
/* Spin-wait loop */
for (uint32_t i = PTHREADPOOL_SPIN_WAIT_ITERATIONS; i != 0; i--) {
pthreadpool_yield();
command = pthreadpool_load_acquire_uint32_t(&threadpool->command);
if (command != last_command) {
return command;
}
}
}
/* Spin-wait disabled or timed out, fall back to mutex/futex wait */
#if PTHREADPOOL_USE_FUTEX
do {
futex_wait(&threadpool->command, last_command);
command = pthreadpool_load_acquire_uint32_t(&threadpool->command);
} while (command == last_command);
#else
/* Lock the command mutex */
pthread_mutex_lock(&threadpool->command_mutex);
/* Read the command */
while ((command = pthreadpool_load_acquire_uint32_t(&threadpool->command)) == last_command) {
/* Wait for new command */
pthread_cond_wait(&threadpool->command_condvar, &threadpool->command_mutex);
}
/* Read a new command */
pthread_mutex_unlock(&threadpool->command_mutex);
#endif
return command;
}
static void* thread_main(void* arg) {
struct thread_info* thread = (struct thread_info*) arg;
struct pthreadpool* threadpool = thread->threadpool;
uint32_t last_command = threadpool_command_init;
struct fpu_state saved_fpu_state = { 0 };
uint32_t flags = 0;
/* Check in */
checkin_worker_thread(threadpool);
/* Monitor new commands and act accordingly */
for (;;) {
uint32_t command = wait_for_new_command(threadpool, last_command, flags);
pthreadpool_fence_acquire();
flags = pthreadpool_load_relaxed_uint32_t(&threadpool->flags);
/* Process command */
switch (command & THREADPOOL_COMMAND_MASK) {
case threadpool_command_parallelize:
{
const thread_function_t thread_function =
(thread_function_t) pthreadpool_load_relaxed_void_p(&threadpool->thread_function);
if (flags & PTHREADPOOL_FLAG_DISABLE_DENORMALS) {
saved_fpu_state = get_fpu_state();
disable_fpu_denormals();
}
thread_function(threadpool, thread);
if (flags & PTHREADPOOL_FLAG_DISABLE_DENORMALS) {
set_fpu_state(saved_fpu_state);
}
break;
}
case threadpool_command_shutdown:
/* Exit immediately: the master thread is waiting on pthread_join */
return NULL;
case threadpool_command_init:
/* To inhibit compiler warning */
break;
}
/* Notify the master thread that we finished processing */
checkin_worker_thread(threadpool);
/* Update last command */
last_command = command;
};
}
struct pthreadpool* pthreadpool_create(size_t threads_count) {
#if PTHREADPOOL_USE_CPUINFO
if (!cpuinfo_initialize()) {
return NULL;
}
#endif
if (threads_count == 0) {
#if PTHREADPOOL_USE_CPUINFO
threads_count = cpuinfo_get_processors_count();
#elif defined(_SC_NPROCESSORS_ONLN)
threads_count = (size_t) sysconf(_SC_NPROCESSORS_ONLN);
#if defined(__EMSCRIPTEN_PTHREADS__)
/* Limit the number of threads to 8 to match link-time PTHREAD_POOL_SIZE option */
if (threads_count >= 8) {
threads_count = 8;
}
#endif
#elif defined(_WIN32)
SYSTEM_INFO system_info;
ZeroMemory(&system_info, sizeof(system_info));
GetSystemInfo(&system_info);
threads_count = (size_t) system_info.dwNumberOfProcessors;
#else
#error "Platform-specific implementation of sysconf(_SC_NPROCESSORS_ONLN) required"
#endif
}
struct pthreadpool* threadpool = pthreadpool_allocate(threads_count);
if (threadpool == NULL) {
return NULL;
}
threadpool->threads_count = fxdiv_init_size_t(threads_count);
for (size_t tid = 0; tid < threads_count; tid++) {
threadpool->threads[tid].thread_number = tid;
threadpool->threads[tid].threadpool = threadpool;
}
/* Thread pool with a single thread computes everything on the caller thread. */
if (threads_count > 1) {
pthread_mutex_init(&threadpool->execution_mutex, NULL);
#if !PTHREADPOOL_USE_FUTEX
pthread_mutex_init(&threadpool->completion_mutex, NULL);
pthread_cond_init(&threadpool->completion_condvar, NULL);
pthread_mutex_init(&threadpool->command_mutex, NULL);
pthread_cond_init(&threadpool->command_condvar, NULL);
#endif
#if PTHREADPOOL_USE_FUTEX
pthreadpool_store_relaxed_uint32_t(&threadpool->has_active_threads, 1);
#endif
pthreadpool_store_relaxed_size_t(&threadpool->active_threads, threads_count - 1 /* caller thread */);
/* Caller thread serves as worker #0. Thus, we create system threads starting with worker #1. */
for (size_t tid = 1; tid < threads_count; tid++) {
pthread_create(&threadpool->threads[tid].thread_object, NULL, &thread_main, &threadpool->threads[tid]);
}
/* Wait until all threads initialize */
wait_worker_threads(threadpool);
}
return threadpool;
}
PTHREADPOOL_INTERNAL void pthreadpool_parallelize(
struct pthreadpool* threadpool,
thread_function_t thread_function,
const void* params,
size_t params_size,
void* task,
void* context,
size_t linear_range,
uint32_t flags)
{
assert(threadpool != NULL);
assert(thread_function != NULL);
assert(task != NULL);
assert(linear_range > 1);
/* Protect the global threadpool structures */
pthread_mutex_lock(&threadpool->execution_mutex);
#if !PTHREADPOOL_USE_FUTEX
/* Lock the command variables to ensure that threads don't start processing before they observe complete command with all arguments */
pthread_mutex_lock(&threadpool->command_mutex);
#endif
/* Setup global arguments */
pthreadpool_store_relaxed_void_p(&threadpool->thread_function, (void*) thread_function);
pthreadpool_store_relaxed_void_p(&threadpool->task, task);
pthreadpool_store_relaxed_void_p(&threadpool->argument, context);
pthreadpool_store_relaxed_uint32_t(&threadpool->flags, flags);
/* Locking of completion_mutex not needed: readers are sleeping on command_condvar */
const struct fxdiv_divisor_size_t threads_count = threadpool->threads_count;
pthreadpool_store_relaxed_size_t(&threadpool->active_threads, threads_count.value - 1 /* caller thread */);
#if PTHREADPOOL_USE_FUTEX
pthreadpool_store_relaxed_uint32_t(&threadpool->has_active_threads, 1);
#endif
if (params_size != 0) {
memcpy(&threadpool->params, params, params_size);
pthreadpool_fence_release();
}
/* Spread the work between threads */
const struct fxdiv_result_size_t range_params = fxdiv_divide_size_t(linear_range, threads_count);
size_t range_start = 0;
for (size_t tid = 0; tid < threads_count.value; tid++) {
struct thread_info* thread = &threadpool->threads[tid];
const size_t range_length = range_params.quotient + (size_t) (tid < range_params.remainder);
const size_t range_end = range_start + range_length;
pthreadpool_store_relaxed_size_t(&thread->range_start, range_start);
pthreadpool_store_relaxed_size_t(&thread->range_end, range_end);
pthreadpool_store_relaxed_size_t(&thread->range_length, range_length);
/* The next subrange starts where the previous ended */
range_start = range_end;
}
/*
* Update the threadpool command.
* Imporantly, do it after initializing command parameters (range, task, argument, flags)
* ~(threadpool->command | THREADPOOL_COMMAND_MASK) flips the bits not in command mask
* to ensure the unmasked command is different then the last command, because worker threads
* monitor for change in the unmasked command.
*/
const uint32_t old_command = pthreadpool_load_relaxed_uint32_t(&threadpool->command);
const uint32_t new_command = ~(old_command | THREADPOOL_COMMAND_MASK) | threadpool_command_parallelize;
/*
* Store the command with release semantics to guarantee that if a worker thread observes
* the new command value, it also observes the updated command parameters.
*
* Note: release semantics is necessary even with a conditional variable, because the workers might
* be waiting in a spin-loop rather than the conditional variable.
*/
pthreadpool_store_release_uint32_t(&threadpool->command, new_command);
#if PTHREADPOOL_USE_FUTEX
/* Wake up the threads */
futex_wake_all(&threadpool->command);
#else
/* Unlock the command variables before waking up the threads for better performance */
pthread_mutex_unlock(&threadpool->command_mutex);
/* Wake up the threads */
pthread_cond_broadcast(&threadpool->command_condvar);
#endif
/* Save and modify FPU denormals control, if needed */
struct fpu_state saved_fpu_state = { 0 };
if (flags & PTHREADPOOL_FLAG_DISABLE_DENORMALS) {
saved_fpu_state = get_fpu_state();
disable_fpu_denormals();
}
/* Do computations as worker #0 */
thread_function(threadpool, &threadpool->threads[0]);
/* Restore FPU denormals control, if needed */
if (flags & PTHREADPOOL_FLAG_DISABLE_DENORMALS) {
set_fpu_state(saved_fpu_state);
}
/* Wait until the threads finish computation */
wait_worker_threads(threadpool);
/* Make changes by other threads visible to this thread */
pthreadpool_fence_acquire();
/* Unprotect the global threadpool structures */
pthread_mutex_unlock(&threadpool->execution_mutex);
}
void pthreadpool_destroy(struct pthreadpool* threadpool) {
if (threadpool != NULL) {
const size_t threads_count = threadpool->threads_count.value;
if (threads_count > 1) {
#if PTHREADPOOL_USE_FUTEX
pthreadpool_store_relaxed_size_t(&threadpool->active_threads, threads_count - 1 /* caller thread */);
pthreadpool_store_relaxed_uint32_t(&threadpool->has_active_threads, 1);
/*
* Store the command with release semantics to guarantee that if a worker thread observes
* the new command value, it also observes the updated active_threads/has_active_threads values.
*/
pthreadpool_store_release_uint32_t(&threadpool->command, threadpool_command_shutdown);
/* Wake up worker threads */
futex_wake_all(&threadpool->command);
#else
/* Lock the command variable to ensure that threads don't shutdown until both command and active_threads are updated */
pthread_mutex_lock(&threadpool->command_mutex);
pthreadpool_store_relaxed_size_t(&threadpool->active_threads, threads_count - 1 /* caller thread */);
/*
* Store the command with release semantics to guarantee that if a worker thread observes
* the new command value, it also observes the updated active_threads value.
*
* Note: the release fence inside pthread_mutex_unlock is insufficient,
* because the workers might be waiting in a spin-loop rather than the conditional variable.
*/
pthreadpool_store_release_uint32_t(&threadpool->command, threadpool_command_shutdown);
/* Wake up worker threads */
pthread_cond_broadcast(&threadpool->command_condvar);
/* Commit the state changes and let workers start processing */
pthread_mutex_unlock(&threadpool->command_mutex);
#endif
/* Wait until all threads return */
for (size_t thread = 1; thread < threads_count; thread++) {
pthread_join(threadpool->threads[thread].thread_object, NULL);
}
/* Release resources */
pthread_mutex_destroy(&threadpool->execution_mutex);
#if !PTHREADPOOL_USE_FUTEX
pthread_mutex_destroy(&threadpool->completion_mutex);
pthread_cond_destroy(&threadpool->completion_condvar);
pthread_mutex_destroy(&threadpool->command_mutex);
pthread_cond_destroy(&threadpool->command_condvar);
#endif
}
#if PTHREADPOOL_USE_CPUINFO
cpuinfo_deinitialize();
#endif
pthreadpool_deallocate(threadpool);
}
}