Persistent pthread thread pool for CK-Engine inference. More...

#include "ck_threadpool.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <errno.h>
#include <unistd.h>
#include <time.h>

Macros
#define	_GNU_SOURCE

#define	CK_SPIN_PAUSE() ((void)0)

Functions
static void	barrier_init (ck_barrier_t *b, int n_threads)

static void	barrier_wait (ck_barrier_t *b)

static int	ck_available_logical_cpus (void)

int	ck_gemm_dynamic_schedule_enabled (void)

int	ck_get_gemm_schedule (void)

int	ck_get_num_threads (void)

int	ck_get_physical_cores (void)

static void	ck_parallel_for_worker (int ith, int nth, void *opaque)

int	ck_set_gemm_schedule (int policy)

void	ck_threadpool_barrier (ck_threadpool_t *pool)

int	ck_threadpool_bounded_capacity (int default_threads, int logical_threads)

int	ck_threadpool_capacity (const ck_threadpool_t *pool)

ck_threadpool_t *	ck_threadpool_create (int n_threads)

ck_threadpool_t *	ck_threadpool_create_capacity (int default_threads, int capacity_threads)

void	ck_threadpool_destroy (ck_threadpool_t *pool)

void	ck_threadpool_dispatch (ck_threadpool_t pool, ck_work_fn_t fn, void args)

void	ck_threadpool_dispatch_n (ck_threadpool_t pool, int active_threads, ck_work_fn_t fn, void args)

ck_threadpool_t *	ck_threadpool_global (void)

void	ck_threadpool_global_destroy (void)

int	ck_threadpool_n_threads (const ck_threadpool_t *pool)

void	ck_threadpool_parallel_for_n (ck_threadpool_t pool, int active_threads, int begin, int end, int grain_size, ck_range_fn_t fn, void args)

void	ck_threadpool_pause (ck_threadpool_t *pool)

void	ck_threadpool_profile_reset (ck_threadpool_t *pool)

void	ck_threadpool_profile_snapshot (const ck_threadpool_t pool, ck_threadpool_profile_t profile)

void	ck_threadpool_resume (ck_threadpool_t *pool)

int	ck_threadpool_thread_id (const ck_threadpool_t *pool)

static void	global_pool_init (void)

static uint64_t	monotonic_ns (void)

static void *	worker_main (void *arg)

Variables
static atomic_int	g_gemm_schedule = CK_GEMM_SCHEDULE_AUTO

static ck_threadpool_t *	g_threadpool = NULL

static pthread_once_t	g_threadpool_once = PTHREAD_ONCE_INIT

Detailed Description

Persistent pthread thread pool for CK-Engine inference.

Architecture:

N-1 worker pthreads created at startup, main thread is thread 0
Workers spin on atomic dispatch counter waiting for work
Barriers use atomic counter + spin-wait with _mm_pause()
Hybrid polling: spin CK_THREADPOOL_SPIN_COUNT rounds, then condvar
All atomics on separate cache lines to avoid false sharing

Based on the ggml_threadpool design from llama.cpp, adapted for CK-Engine's kernel dispatch model.

Definition in file ck_threadpool.c.

Macro Definition Documentation

◆ _GNU_SOURCE

#define _GNU_SOURCE

Definition at line 2 of file ck_threadpool.c.

◆ CK_SPIN_PAUSE

#define CK_SPIN_PAUSE ( ) ((void)0)

Definition at line 36 of file ck_threadpool.c.

Function Documentation

◆ barrier_init()

static void barrier_init	(	ck_barrier_t *	b,
		int	n_threads
	)

static

Definition at line 124 of file ck_threadpool.c.

{
    atomic_store(&b->n_arrived, 0);
    atomic_store(&b->n_phase, 0);
    b->n_threads = n_threads;
}

Referenced by ck_threadpool_create_capacity(), and ck_threadpool_dispatch_n().

◆ barrier_wait()

static void barrier_wait ( ck_barrier_t * b )

static

Spin-wait barrier. All threads must call this. Uses phase counter to allow re-use without reset.

Definition at line 135 of file ck_threadpool.c.

{
    const int n = b->n_threads;
    const int phase = atomic_load_explicit(&b->n_phase, memory_order_relaxed);
 
    if (atomic_fetch_add_explicit(&b->n_arrived, 1, memory_order_acq_rel) == n - 1) {
        /* Last thread to arrive — reset and advance phase */
        atomic_store_explicit(&b->n_arrived, 0, memory_order_relaxed);
        atomic_store_explicit(&b->n_phase, phase + 1, memory_order_release);
    } else {
        /* Spin until phase advances */
        int spins = 0;
        while (atomic_load_explicit(&b->n_phase, memory_order_acquire) == phase) {
            CK_SPIN_PAUSE();
            spins++;
            /* After many spins, yield to avoid wasting CPU on oversubscribed systems */
            if (spins > CK_THREADPOOL_SPIN_COUNT * 16) {
                sched_yield();
                spins = 0;
            }
        }
    }
}

References CK_SPIN_PAUSE, and CK_THREADPOOL_SPIN_COUNT.

Referenced by ck_threadpool_barrier().

◆ ck_available_logical_cpus()

static int ck_available_logical_cpus ( void )

static

Definition at line 608 of file ck_threadpool.c.

{
#ifdef __linux__
    cpu_set_t allowed;
    if (sched_getaffinity(0, sizeof(allowed), &allowed) == 0) {
        int count = 0;
        for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) {
            if (CPU_ISSET(cpu, &allowed)) ++count;
        }
        if (count > 0) return count;
    }
#endif
    const long online = sysconf(_SC_NPROCESSORS_ONLN);
    return online > 0 ? (int)online : 1;
}

Referenced by global_pool_init().

◆ ck_gemm_dynamic_schedule_enabled()

int ck_gemm_dynamic_schedule_enabled ( void )

Return non-zero when independent GEMM tiles should use dynamic claiming.

Definition at line 55 of file ck_threadpool.c.

{
    const int policy = ck_get_gemm_schedule();
    return policy == CK_GEMM_SCHEDULE_AUTO || policy == CK_GEMM_SCHEDULE_DYNAMIC;
}

References CK_GEMM_SCHEDULE_AUTO, CK_GEMM_SCHEDULE_DYNAMIC, and ck_get_gemm_schedule().

Referenced by gemm_nt_q4_k_packed_vnni_x8_q8_k_split_min_threaded_4m().

◆ ck_get_gemm_schedule()

int ck_get_gemm_schedule ( void )

Return the configured process-wide GEMM scheduling policy.

Definition at line 50 of file ck_threadpool.c.

{
    return atomic_load_explicit(&g_gemm_schedule, memory_order_acquire);
}

References g_gemm_schedule.

Referenced by ck_gemm_dynamic_schedule_enabled().

◆ ck_get_num_threads()

int ck_get_num_threads ( void )

extern

Definition at line 310 of file ckernel_strict.c.

{
    // Auto-initialize if not set
    if (!g_threads_initialized) {
        ck_set_num_threads(0);  // Auto-detect
    }
    return g_num_threads;
}

Referenced by global_pool_init().

◆ ck_get_physical_cores()

int ck_get_physical_cores ( void )

extern

Definition at line 170 of file ckernel_strict.c.

{
    int physical_cores = 0;
    int logical_cores = (int)sysconf(_SC_NPROCESSORS_ONLN);
    if (logical_cores <= 0) {
        logical_cores = 1;
    }
 
    int cpu_cores_hint = 0;
    int siblings_hint = 0;
 
    // Read from /proc/cpuinfo (Linux) and count unique (physical id, core id) pairs.
    FILE *f = fopen("/proc/cpuinfo", "r");
    if (f) {
        char line[256];
        int physical_id = -1;
        int core_id = -1;
 
        struct {
            int physical_id;
            int core_id;
        } seen[8192];
        int seen_count = 0;
 
        const int seen_cap = (int)(sizeof(seen) / sizeof(seen[0]));
 
        // Helper: add (pid,cid) to set if not present.
        #define CK_ADD_PAIR(pid, cid)                                            \
            do {                                                                 \
                if ((pid) >= 0 && (cid) >= 0) {                                  \
                    int exists = 0;                                              \
                    for (int ii = 0; ii < seen_count; ++ii) {                    \
                        if (seen[ii].physical_id == (pid) &&                     \
                            seen[ii].core_id == (cid)) {                         \
                            exists = 1;                                          \
                            break;                                               \
                        }                                                        \
                    }                                                            \
                    if (!exists && seen_count < seen_cap) {                      \
                        seen[seen_count].physical_id = (pid);                    \
                        seen[seen_count].core_id = (cid);                        \
                        ++seen_count;                                            \
                    }                                                            \
                }                                                                \
            } while (0)
 
        while (fgets(line, sizeof(line), f)) {
            int val;
 
            // Blank line separates processor blocks.
            if (line[0] == '\n' || line[0] == '\0') {
                CK_ADD_PAIR(physical_id, core_id);
                physical_id = -1;
                core_id = -1;
                continue;
            }
 
            if (sscanf(line, "physical id : %d", &val) == 1) {
                physical_id = val;
                continue;
            }
            if (sscanf(line, "core id : %d", &val) == 1) {
                core_id = val;
                continue;
            }
            if (sscanf(line, "cpu cores : %d", &val) == 1) {
                if (val > cpu_cores_hint) cpu_cores_hint = val;
                continue;
            }
            if (sscanf(line, "siblings : %d", &val) == 1) {
                if (val > siblings_hint) siblings_hint = val;
                continue;
            }
        }
        fclose(f);
 
        // Handle file without trailing blank line.
        CK_ADD_PAIR(physical_id, core_id);
 
        #undef CK_ADD_PAIR
 
        physical_cores = seen_count;
    }
 
    // Fallback: infer threads-per-core from siblings/cpu cores when pair data
    // is missing (common in containers/VMs).
    if (physical_cores <= 1 && logical_cores > 1) {
        int threads_per_core = 0;
        if (siblings_hint > 0 && cpu_cores_hint > 0 && siblings_hint >= cpu_cores_hint) {
            threads_per_core = siblings_hint / cpu_cores_hint;
        }
        if (threads_per_core > 1) {
            int inferred_physical = logical_cores / threads_per_core;
            if (inferred_physical > 1) {
                return inferred_physical;
            }
        }
        if (cpu_cores_hint > 1 && cpu_cores_hint <= logical_cores) {
            return cpu_cores_hint;
        }
        return logical_cores;
    }
 
    if (physical_cores > 1) {
        return physical_cores;
    }
 
    return logical_cores;
}

Referenced by ck_set_num_threads(), ck_threadpool_create_capacity(), and global_pool_init().

◆ ck_parallel_for_worker()

static void ck_parallel_for_worker	(	int	ith,
		int	nth,
		void *	opaque
	)

static

Definition at line 479 of file ck_threadpool.c.

{
    (void)ith;
    (void)nth;
    ck_parallel_for_work_t *work = (ck_parallel_for_work_t *)opaque;
    for (;;) {
        const int begin = atomic_fetch_add_explicit(
            &work->next, work->grain_size, memory_order_relaxed);
        if (begin >= work->end) break;
        int end = begin + work->grain_size;
        if (end > work->end) end = work->end;
        work->fn(begin, end, work->args);
    }
}

References end.

Referenced by ck_threadpool_parallel_for_n().

◆ ck_set_gemm_schedule()

int ck_set_gemm_schedule ( int policy )

Set the process-wide GEMM tile scheduling policy.

AUTO is the production default and currently selects dynamic work claiming for providers whose jobs write independent output tiles. Providers with ordered/shared reductions do not consult this policy.

Returns: 0 on success, -1 for an invalid policy.

Definition at line 41 of file ck_threadpool.c.

{
    if (policy < CK_GEMM_SCHEDULE_AUTO || policy > CK_GEMM_SCHEDULE_DYNAMIC) {
        return -1;
    }
    atomic_store_explicit(&g_gemm_schedule, policy, memory_order_release);
    return 0;
}

References CK_GEMM_SCHEDULE_DYNAMIC, and g_gemm_schedule.

◆ ck_threadpool_barrier()

void ck_threadpool_barrier ( ck_threadpool_t * pool )

Barrier synchronization within a dispatched work function.

ALL threads must call this at the same point. Threads spin-wait until all have arrived, then proceed.

Must only be called from within a work function (during dispatch).

Parameters

pool	Thread pool

Definition at line 520 of file ck_threadpool.c.

{
    if (!pool || pool->n_threads <= 1) return;
    barrier_wait(&pool->barrier);
}

References barrier_wait().

◆ ck_threadpool_bounded_capacity()

int ck_threadpool_bounded_capacity	(	int	default_threads,
		int	logical_threads
	)

Compute the bounded capacity for an SMT-safe provider. The default width is preserved and at most half of the additional logical CPUs are reserved.

Definition at line 256 of file ck_threadpool.c.

{
    if (default_threads < 1) default_threads = 1;
    if (default_threads > CK_THREADPOOL_MAX_THREADS) {
        default_threads = CK_THREADPOOL_MAX_THREADS;
    }
    if (logical_threads <= default_threads) return default_threads;
 
    int capacity = default_threads + (logical_threads - default_threads) / 2;
    if (capacity > CK_THREADPOOL_MAX_THREADS) {
        capacity = CK_THREADPOOL_MAX_THREADS;
    }
    return capacity;
}

References CK_THREADPOOL_MAX_THREADS.

Referenced by global_pool_init().

◆ ck_threadpool_capacity()

int ck_threadpool_capacity ( const ck_threadpool_t * pool )

Get the maximum worker capacity available to explicit dispatch_n calls.

Definition at line 556 of file ck_threadpool.c.

{
    return pool ? pool->n_threads : 1;
}

Referenced by gemm_nt_q4_k_packed_vnni_x16_q8_k_split_min_threaded_16m(), and gemm_nt_q4_k_packed_vnni_x8_q8_k_split_min_threaded_4m().

◆ ck_threadpool_create()

ck_threadpool_t * ck_threadpool_create ( int n_threads )

Create a thread pool with n_threads total threads. Thread 0 is the calling (main) thread; n_threads-1 workers are spawned.

Parameters

n_threads Total thread count (including main). Must be >= 1. Pass 0 for auto-detect (physical cores).

Returns: Pool handle, or NULL on failure.

Definition at line 342 of file ck_threadpool.c.

{
    return ck_threadpool_create_capacity(n_threads, n_threads);
}

References ck_threadpool_create_capacity().

◆ ck_threadpool_create_capacity()

ck_threadpool_t * ck_threadpool_create_capacity	(	int	default_threads,
		int	capacity_threads
	)

Create a pool whose ordinary dispatch width is smaller than its worker capacity. Exact providers may opt into the additional workers with ck_threadpool_dispatch_n(); ordinary dispatch remains at default_threads.

Definition at line 271 of file ck_threadpool.c.

{
    if (default_threads <= 0) {
        default_threads = ck_get_physical_cores();
        if (default_threads <= 0) default_threads = 1;
        /* Cap at reasonable default for memory-bound workloads */
        if (default_threads > 8) default_threads = 8;
    }
    if (capacity_threads < default_threads) {
        capacity_threads = default_threads;
    }
    if (capacity_threads > CK_THREADPOOL_MAX_THREADS) {
        capacity_threads = CK_THREADPOOL_MAX_THREADS;
    }
    if (default_threads > capacity_threads) {
        default_threads = capacity_threads;
    }
 
    ck_threadpool_t *pool = aligned_alloc(CK_CACHE_LINE, sizeof(ck_threadpool_t));
    if (!pool) return NULL;
    memset(pool, 0, sizeof(*pool));
 
    pool->n_threads = capacity_threads;
    pool->default_threads = default_threads;
    atomic_store(&pool->n_dispatch, 0);
    atomic_store(&pool->n_complete, 0);
    atomic_store(&pool->active_threads, default_threads);
    atomic_store(&pool->stop, 0);
    atomic_store(&pool->paused, 0);
    atomic_store(&pool->profile_enabled, 0);
    pool->work_fn = NULL;
    pool->work_args = NULL;
 
    barrier_init(&pool->barrier, default_threads);
 
    pthread_mutex_init(&pool->mutex, NULL);
    pthread_cond_init(&pool->cond_dispatch, NULL);
    pthread_cond_init(&pool->cond_done, NULL);
 
    /* Thread 0 = main thread (no pthread created) */
    pool->workers[0].id = 0;
    pool->workers[0].pool = pool;
    pool->workers[0].thread = pthread_self();
 
    /* Spawn N-1 worker threads */
    for (int i = 1; i < capacity_threads; i++) {
        pool->workers[i].id = i;
        pool->workers[i].pool = pool;
 
        int rc = pthread_create(&pool->workers[i].thread, NULL,
                                worker_main, &pool->workers[i]);
        if (rc != 0) {
            fprintf(stderr, "[CK threadpool] Failed to create worker %d: %s\n",
                    i, strerror(rc));
            /* Reduce thread count to what we managed to create */
            pool->n_threads = i;
            barrier_init(&pool->barrier, i);
            break;
        }
    }
 
    if (pool->n_threads > 1) {
        fprintf(stderr,
                "[CK threadpool] Created %d threads (default=%d, 1 main + %d workers)\n",
                pool->n_threads, pool->default_threads, pool->n_threads - 1);
    }
 
    return pool;
}

References barrier_init(), CK_CACHE_LINE, ck_get_physical_cores(), CK_THREADPOOL_MAX_THREADS, and worker_main().

Referenced by ck_threadpool_create(), and global_pool_init().

◆ ck_threadpool_destroy()

void ck_threadpool_destroy ( ck_threadpool_t * pool )

Destroy the thread pool. Signals all workers to exit and joins them. Safe to call with NULL.

Definition at line 347 of file ck_threadpool.c.

{
    if (!pool) return;
 
    /* Signal shutdown */
    atomic_store_explicit(&pool->stop, 1, memory_order_release);
 
    /* Wake all sleeping workers */
    pthread_mutex_lock(&pool->mutex);
    pthread_cond_broadcast(&pool->cond_dispatch);
    pthread_mutex_unlock(&pool->mutex);
 
    /* Join all worker threads */
    for (int i = 1; i < pool->n_threads; i++) {
        pthread_join(pool->workers[i].thread, NULL);
    }
 
    pthread_cond_destroy(&pool->cond_dispatch);
    pthread_cond_destroy(&pool->cond_done);
    pthread_mutex_destroy(&pool->mutex);
 
    free(pool);
}

Referenced by ck_threadpool_global_destroy().

◆ ck_threadpool_dispatch()

void ck_threadpool_dispatch	(	ck_threadpool_t *	pool,
		ck_work_fn_t	fn,
		void *	args
	)

Dispatch work to all threads and wait for completion.

Sets the work function and args
Bumps the dispatch counter (wakes workers)
Main thread (ith=0) executes its share
Waits for all threads to complete via barrier

This is a blocking call — returns when ALL threads have finished.

Parameters

pool	Thread pool
fn	Work function (called on each thread)
args	Argument passed to fn

Definition at line 465 of file ck_threadpool.c.

{
    if (!pool) return;
    ck_threadpool_dispatch_n(pool, pool->default_threads, fn, args);
}

References ck_threadpool_dispatch_n().

◆ ck_threadpool_dispatch_n()

void ck_threadpool_dispatch_n	(	ck_threadpool_t *	pool,
		int	active_threads,
		ck_work_fn_t	fn,
		void *	args
	)

Dispatch work to a subset of the pool and wait for completion.

Threads with ith >= active_threads remain idle for this dispatch. The work function sees nth == active_threads.

Parameters

pool	Thread pool
active_threads	Number of active threads including main thread
fn	Work function
args	Argument passed to fn

Definition at line 375 of file ck_threadpool.c.

{
    if (!pool || !fn) return;
    if (active_threads <= 0) {
        active_threads = 1;
    }
    if (active_threads > pool->n_threads) {
        active_threads = pool->n_threads;
    }
 
    const int profile = atomic_load_explicit(
        &pool->profile_enabled, memory_order_relaxed);
    const uint64_t dispatch_start = profile ? monotonic_ns() : 0;
 
    /* Single-thread fast path: just call directly */
    if (active_threads == 1 || pool->n_threads == 1) {
        fn(0, 1, args);
        if (profile) {
            const uint64_t dispatch_end = monotonic_ns();
            atomic_fetch_add_explicit(&pool->profile_dispatch_count, 1, memory_order_relaxed);
            atomic_fetch_add_explicit(
                &pool->profile_dispatch_total_ns,
                dispatch_end - dispatch_start,
                memory_order_relaxed);
            atomic_fetch_add_explicit(
                &pool->profile_main_work_ns,
                dispatch_end - dispatch_start,
                memory_order_relaxed);
        }
        return;
    }
 
    /* Reset barrier phase for this dispatch */
    barrier_init(&pool->barrier, active_threads);
 
    /* Set work descriptor */
    pthread_mutex_lock(&pool->mutex);
    pool->work_fn = fn;
    pool->work_args = args;
    atomic_store_explicit(&pool->active_threads, active_threads, memory_order_release);
    atomic_store_explicit(&pool->n_complete, 0, memory_order_release);
 
    /* Wake workers by bumping dispatch counter */
    atomic_fetch_add_explicit(&pool->n_dispatch, 1, memory_order_release);
 
    /* Also signal condvar for sleeping workers */
    pthread_cond_broadcast(&pool->cond_dispatch);
    pthread_mutex_unlock(&pool->mutex);
 
    /* Main thread (ith=0) does its share */
    const uint64_t main_start = profile ? monotonic_ns() : 0;
    fn(0, active_threads, args);
    const uint64_t main_end = profile ? monotonic_ns() : 0;
 
    /* Wait for all workers to complete */
    if (active_threads > 1) {
        int spins = 0;
        while (atomic_load_explicit(&pool->n_complete, memory_order_acquire)
               < active_threads - 1) {
            CK_SPIN_PAUSE();
            spins++;
            if (spins >= CK_THREADPOOL_SPIN_COUNT) {
                pthread_mutex_lock(&pool->mutex);
                if (atomic_load_explicit(&pool->n_complete, memory_order_acquire)
                    < active_threads - 1) {
                    pthread_cond_wait(&pool->cond_done, &pool->mutex);
                }
                pthread_mutex_unlock(&pool->mutex);
                spins = 0;
            }
        }
    }
    if (profile) {
        const uint64_t dispatch_end = monotonic_ns();
        atomic_fetch_add_explicit(&pool->profile_dispatch_count, 1, memory_order_relaxed);
        atomic_fetch_add_explicit(
            &pool->profile_dispatch_total_ns,
            dispatch_end - dispatch_start,
            memory_order_relaxed);
        atomic_fetch_add_explicit(
            &pool->profile_main_work_ns,
            main_end - main_start,
            memory_order_relaxed);
        atomic_fetch_add_explicit(
            &pool->profile_completion_wait_ns,
            dispatch_end - main_end,
            memory_order_relaxed);
    }
}

References barrier_init(), CK_SPIN_PAUSE, CK_THREADPOOL_SPIN_COUNT, and monotonic_ns().

◆ ck_threadpool_global()

ck_threadpool_t * ck_threadpool_global ( void )

Get or create the global thread pool. Thread-safe (uses pthread_once internally). Uses ck_get_num_threads() for the default width. In automatic mode the pool may reserve a bounded subset of SMT siblings as explicit provider capacity; ordinary dispatch remains at the default width.

Returns: Global pool, never NULL after successful first call.

Definition at line 651 of file ck_threadpool.c.

{
    pthread_once(&g_threadpool_once, global_pool_init);
    return g_threadpool;
}

References g_threadpool, g_threadpool_once, and global_pool_init().

◆ ck_threadpool_global_destroy()

void ck_threadpool_global_destroy ( void )

Destroy the global thread pool. Called during engine shutdown.

Definition at line 657 of file ck_threadpool.c.

{
    if (g_threadpool) {
        ck_threadpool_destroy(g_threadpool);
        g_threadpool = NULL;
        /* Reset once control so pool can be re-created if needed */
        g_threadpool_once = PTHREAD_ONCE_INIT;
    }
}

References ck_threadpool_destroy(), g_threadpool, and g_threadpool_once.

Referenced by ck_threadpool_shutdown().

◆ ck_threadpool_n_threads()

int ck_threadpool_n_threads ( const ck_threadpool_t * pool )

Get the ordinary/default dispatch width (including the main thread).

Definition at line 551 of file ck_threadpool.c.

{
    return pool ? pool->default_threads : 1;
}

◆ ck_threadpool_parallel_for_n()

void ck_threadpool_parallel_for_n	(	ck_threadpool_t *	pool,
		int	active_threads,
		int	begin,
		int	end,
		int	grain_size,
		ck_range_fn_t	fn,
		void *	args
	)

Dynamically distribute independent ranges through the persistent pool.

Workers claim grain_size consecutive indices until [begin, end) is empty. This changes ownership only; callers remain responsible for ensuring that ranges write disjoint outputs and preserve each output's reduction order. Do not use this helper for unordered shared reductions.

Definition at line 494 of file ck_threadpool.c.

{
    if (!fn || begin >= end) return;
    if (grain_size <= 0) grain_size = 1;
    if (!pool || active_threads <= 1) {
        fn(begin, end, args);
        return;
    }
 
    ck_parallel_for_work_t work = {
        .end = end,
        .grain_size = grain_size,
        .fn = fn,
        .args = args,
    };
    atomic_init(&work.next, begin);
    ck_threadpool_dispatch_n(
        pool, active_threads, ck_parallel_for_worker, &work);
}

References ck_parallel_for_worker(), ck_threadpool_dispatch_n(), and end.

Referenced by audio_conv2d_whc_grouped_f32(), audio_glu_split_channel_major_f32(), audio_relative_shift_f32(), ck_gemma4_prepare_parallel(), deepseek_mla_kv_decompress_bf16_parallel_dispatch(), gemm_nt_bf16_bf16_storage_parallel_dispatch(), gemm_nt_bf16_parallel_dispatch(), gemm_nt_fp32_exact_parallel_dispatch(), gemm_nt_q4_k_packed_vnni_x8_q8_k_split_min_threaded_4m(), gemma4_per_layer_embed_forward(), gemv_bf16_bf16_storage_parallel_dispatch(), gemv_bf16_parallel_dispatch(), gemv_nvfp4_q8_0_uniform(), ssm_conv1d_forward_llama_fma(), and ssm_conv1d_forward_llama_production().

◆ ck_threadpool_pause()

void ck_threadpool_pause ( ck_threadpool_t * pool )

Pause workers — they sleep on condvar (0% CPU). Call between batches or during interactive waiting. Workers wake on next dispatch or resume.

Definition at line 530 of file ck_threadpool.c.

{
    if (!pool) return;
    atomic_store_explicit(&pool->paused, 1, memory_order_release);
}

◆ ck_threadpool_profile_reset()

void ck_threadpool_profile_reset ( ck_threadpool_t * pool )

Enable profiling and reset cumulative dispatch timing counters.

Definition at line 573 of file ck_threadpool.c.

{
    if (!pool) return;
    atomic_store_explicit(&pool->profile_dispatch_count, 0, memory_order_relaxed);
    atomic_store_explicit(&pool->profile_dispatch_total_ns, 0, memory_order_relaxed);
    atomic_store_explicit(&pool->profile_main_work_ns, 0, memory_order_relaxed);
    atomic_store_explicit(&pool->profile_completion_wait_ns, 0, memory_order_relaxed);
    atomic_store_explicit(&pool->profile_enabled, 1, memory_order_release);
}

◆ ck_threadpool_profile_snapshot()

void ck_threadpool_profile_snapshot	(	const ck_threadpool_t *	pool,
		ck_threadpool_profile_t *	profile
	)

Snapshot cumulative dispatch timing counters without stopping workers.

Definition at line 583 of file ck_threadpool.c.

{
    if (!profile) return;
    memset(profile, 0, sizeof(*profile));
    if (!pool) return;
    profile->dispatch_count = atomic_load_explicit(
        &pool->profile_dispatch_count, memory_order_relaxed);
    profile->dispatch_total_ns = atomic_load_explicit(
        &pool->profile_dispatch_total_ns, memory_order_relaxed);
    profile->main_work_ns = atomic_load_explicit(
        &pool->profile_main_work_ns, memory_order_relaxed);
    profile->completion_wait_ns = atomic_load_explicit(
        &pool->profile_completion_wait_ns, memory_order_relaxed);
}

References ck_threadpool_profile_t::completion_wait_ns, ck_threadpool_profile_t::dispatch_count, ck_threadpool_profile_t::dispatch_total_ns, and ck_threadpool_profile_t::main_work_ns.

◆ ck_threadpool_resume()

void ck_threadpool_resume ( ck_threadpool_t * pool )

Resume workers — transition from sleep to spin-wait. Call before starting a new batch of work.

Definition at line 536 of file ck_threadpool.c.

{
    if (!pool) return;
    atomic_store_explicit(&pool->paused, 0, memory_order_release);
 
    /* Wake sleeping workers */
    pthread_mutex_lock(&pool->mutex);
    pthread_cond_broadcast(&pool->cond_dispatch);
    pthread_mutex_unlock(&pool->mutex);
}

◆ ck_threadpool_thread_id()

int ck_threadpool_thread_id ( const ck_threadpool_t * pool )

Get thread index for current thread (0 = main, -1 if not in pool)

Definition at line 561 of file ck_threadpool.c.

{
    if (!pool) return -1;
    pthread_t self = pthread_self();
    for (int i = 0; i < pool->n_threads; i++) {
        if (pthread_equal(self, pool->workers[i].thread)) {
            return i;
        }
    }
    return -1;
}

Referenced by attention_forward_causal_head_major_gqa_prefill_append_bf16cache_pytorch_contract_workspace(), attention_forward_causal_head_major_gqa_prefill_append_f16cache_gqa_reuse_config(), ck_attention_f16_prefill_qtile64_dispatch(), ck_attention_forward_causal_head_major_gqa_prefill_segmented_f16cache_schedule_workspace(), and ck_sliding_attention_pick_threads().

◆ global_pool_init()

static void global_pool_init ( void )

static

Definition at line 624 of file ck_threadpool.c.

{
    const int available_threads = ck_available_logical_cpus();
    int physical_threads = ck_get_physical_cores();
    if (physical_threads > available_threads) physical_threads = available_threads;
    int default_threads = ck_get_num_threads();
    if (default_threads > available_threads) {
        default_threads = available_threads;
    }
    int capacity_threads = default_threads;
    const char *capacity_env = getenv("CK_THREADPOOL_CAPACITY");
    if (capacity_env && atoi(capacity_env) > 0) {
        capacity_threads = atoi(capacity_env);
    } else if (!getenv("CK_NUM_THREADS") &&
               default_threads == physical_threads) {
        /* Generated runtimes set OMP_NUM_THREADS=1 to keep OpenMP dormant and
         * configure the CK pool separately. Do not mistake that isolation
         * setting for a CK capacity cap. A non-physical default remains an
         * explicit width and does not gain automatic SMT workers. */
        capacity_threads = ck_threadpool_bounded_capacity(
                default_threads, available_threads);
    }
    if (capacity_threads > available_threads) capacity_threads = available_threads;
    g_threadpool = ck_threadpool_create_capacity(
            default_threads, capacity_threads);
}

References ck_available_logical_cpus(), ck_get_num_threads(), ck_get_physical_cores(), ck_threadpool_bounded_capacity(), ck_threadpool_create_capacity(), and g_threadpool.

Referenced by ck_threadpool_global().

◆ monotonic_ns()

static uint64_t monotonic_ns ( void )

static

Definition at line 113 of file ck_threadpool.c.

{
    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC, &now);
    return (uint64_t)now.tv_sec * UINT64_C(1000000000) + (uint64_t)now.tv_nsec;
}

Referenced by ck_threadpool_dispatch_n().

◆ worker_main()

static void * worker_main ( void * arg )

static

Definition at line 163 of file ck_threadpool.c.

{
    ck_worker_t *w = (ck_worker_t *)arg;
    ck_threadpool_t *pool = w->pool;
    const int ith = w->id;
    int last_dispatch = 0;
 
    for (;;) {
        /* Spin-wait for new dispatch */
        int spins = 0;
        int active = 0;
        ck_work_fn_t fn = NULL;
        void *args = NULL;
        for (;;) {
            /* Check shutdown */
            if (atomic_load_explicit(&pool->stop, memory_order_acquire)) {
                return NULL;
            }
 
            /* Check for new work */
            int current = atomic_load_explicit(&pool->n_dispatch, memory_order_acquire);
            active = atomic_load_explicit(&pool->active_threads, memory_order_acquire);
            if (current != last_dispatch) {
                /* Snapshot the epoch and descriptor together. An inactive
                 * worker may still be catching up with an older dispatch. */
                pthread_mutex_lock(&pool->mutex);
                current = atomic_load_explicit(&pool->n_dispatch, memory_order_acquire);
                active = atomic_load_explicit(&pool->active_threads, memory_order_acquire);
                if (current != last_dispatch) {
                    last_dispatch = current;
                    if (ith < active) {
                        fn = pool->work_fn;
                        args = pool->work_args;
                        pthread_mutex_unlock(&pool->mutex);
                        break;
                    }
                }
                pthread_mutex_unlock(&pool->mutex);
                spins = 0;
            }
 
            /* Threads outside the active subset sleep instead of spinning. */
            if (ith >= active || spins >= CK_THREADPOOL_SPIN_COUNT) {
                pthread_mutex_lock(&pool->mutex);
                for (;;) {
                    if (atomic_load_explicit(&pool->stop, memory_order_acquire)) {
                        pthread_mutex_unlock(&pool->mutex);
                        return NULL;
                    }
                    current = atomic_load_explicit(&pool->n_dispatch, memory_order_acquire);
                    active = atomic_load_explicit(&pool->active_threads, memory_order_acquire);
                    if (current != last_dispatch) {
                        last_dispatch = current;
                        if (ith < active) {
                            fn = pool->work_fn;
                            args = pool->work_args;
                            pthread_mutex_unlock(&pool->mutex);
                            goto worker_have_work;
                        }
                    }
                    pthread_cond_wait(&pool->cond_dispatch, &pool->mutex);
                }
            }
 
            CK_SPIN_PAUSE();
            spins++;
        }
 
worker_have_work:
        /* Execute work */
        if (fn) {
            fn(ith, active, args);
        }
 
        /* Signal completion */
        if (atomic_fetch_add_explicit(&pool->n_complete, 1, memory_order_acq_rel)
            == active - 2) {
            /* Last worker done — wake main thread if it's waiting */
            pthread_mutex_lock(&pool->mutex);
            pthread_cond_signal(&pool->cond_done);
            pthread_mutex_unlock(&pool->mutex);
        }
    }
 
    return NULL;
}

References CK_SPIN_PAUSE, and CK_THREADPOOL_SPIN_COUNT.

Referenced by ck_threadpool_create_capacity().

Variable Documentation

◆ g_gemm_schedule

atomic_int g_gemm_schedule = CK_GEMM_SCHEDULE_AUTO

static

Definition at line 39 of file ck_threadpool.c.

Referenced by ck_get_gemm_schedule(), and ck_set_gemm_schedule().

◆ g_threadpool

ck_threadpool_t* g_threadpool = NULL

static

Definition at line 603 of file ck_threadpool.c.

Referenced by ck_threadpool_global(), ck_threadpool_global_destroy(), and global_pool_init().

◆ g_threadpool_once

pthread_once_t g_threadpool_once = PTHREAD_ONCE_INIT

static

Definition at line 604 of file ck_threadpool.c.

Referenced by ck_threadpool_global(), and ck_threadpool_global_destroy().

Macros

Functions

Variables

Detailed Description

Macro Definition Documentation

◆ _GNU_SOURCE

◆ CK_SPIN_PAUSE

Function Documentation

◆ barrier_init()

◆ barrier_wait()

◆ ck_available_logical_cpus()

◆ ck_gemm_dynamic_schedule_enabled()

◆ ck_get_gemm_schedule()

◆ ck_get_num_threads()

◆ ck_get_physical_cores()

◆ ck_parallel_for_worker()

◆ ck_set_gemm_schedule()

◆ ck_threadpool_barrier()

◆ ck_threadpool_bounded_capacity()

◆ ck_threadpool_capacity()

◆ ck_threadpool_create()

◆ ck_threadpool_create_capacity()

◆ ck_threadpool_destroy()

◆ ck_threadpool_dispatch()

◆ ck_threadpool_dispatch_n()

◆ ck_threadpool_global()

◆ ck_threadpool_global_destroy()

◆ ck_threadpool_n_threads()

◆ ck_threadpool_parallel_for_n()

◆ ck_threadpool_pause()

◆ ck_threadpool_profile_reset()

◆ ck_threadpool_profile_snapshot()

◆ ck_threadpool_resume()

◆ ck_threadpool_thread_id()

◆ global_pool_init()

◆ monotonic_ns()

◆ worker_main()

Variable Documentation

◆ g_gemm_schedule

◆ g_threadpool

◆ g_threadpool_once