system-info: detect thread/L1 sizes at runtime, drop build.rs

Author: TomWambsgans

crates/backend/parallel/src/lib.rs

@@ -23,8 +23,6 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::{Mutex, Once, OnceLock};
 use std::thread::Thread;
 
-use system_info::NUM_THREADS;
-
 /// Idle spins before a worker parks: long enough to stay hot across back-to-back dispatches,
 /// short enough to yield the core during sequential gaps.
 const SPIN_LIMIT: u32 = 1 << 12;
@@ -33,18 +31,18 @@ const SPIN_LIMIT: u32 = 1 << 12;
 /// million-task kernels to a few thousand claims.
 const MAX_CLAIM_BATCH: usize = 1 << 12;
 
-/// Worker count including the dispatcher (= build-time `NUM_THREADS`).
+/// Worker count including the dispatcher. Resolved once at runtime (see [`system_info::num_threads`]).
 #[must_use]
-pub const fn num_threads() -> usize {
-    NUM_THREADS
+pub fn num_threads() -> usize {
+    system_info::num_threads()
 }
 
 /// Chunk size for a flat fan-out: a few chunks per worker — fine enough for the counter to
 /// rebalance heterogeneous cores, coarse enough to amortize dispatch.
 #[must_use]
 #[inline]
 pub fn recommended_chunk_size(n_items: usize) -> usize {
-    n_items.div_ceil(NUM_THREADS * 4).max(1)
+    n_items.div_ceil(num_threads() * 4).max(1)
 }
 
 thread_local! {
@@ -107,28 +105,20 @@ unsafe impl Send for Pool {}
 
 /// Idempotent warm-up: spawn workers and run one empty dispatch so the pool and the (macOS)
 /// lazily-allocated mutex exist before timed work; otherwise the pool inits on first use.
-///
-/// Also fail-fast if the machine's core count differs from the build-time [`NUM_THREADS`] (which
-/// sizes the pool): a mismatch silently over/under-subscribes every kernel.
 pub fn init() {
     static INIT: Once = Once::new();
     INIT.call_once(|| {
-        let actual = std::thread::available_parallelism().unwrap().get();
-        assert_eq!(
-            actual, NUM_THREADS,
-            "parallel pool built for {NUM_THREADS} threads but this machine reports {actual} -> please rebuild with env variable: LEANVM_NUM_THREADS={actual}"
-        );
         let _ = pool();
-        if NUM_THREADS > 1 {
-            for_each_index(NUM_THREADS, |_| {});
+        if num_threads() > 1 {
+            for_each_index(num_threads(), |_| {});
         }
     });
 }
 
 fn pool() -> &'static Pool {
     static POOL: OnceLock<&'static Pool> = OnceLock::new();
     POOL.get_or_init(|| {
-        let n = NUM_THREADS.max(1);
+        let n = num_threads().max(1);
         let p: &'static Pool = Box::leak(Box::new(Pool {
             job: UnsafeCell::new(None),
             generation: AtomicUsize::new(0),
@@ -200,6 +190,7 @@ fn drain(pool: &Pool) {
     // SAFETY: `job.f` borrows a `&dyn Fn` the blocked dispatcher keeps live.
     let f = unsafe { job.f.as_ref() };
     let n = job.n_tasks;
+    let nt = num_threads();
     let prev = IN_TASK.replace(true); // catch nested dispatch (see `for_each_chunk`)
     // Catch a task panic so it can't unwind across `worker_main` (skipping the `working`
     // decrement → deadlock) or poison the dispatch lock; `for_each_chunk` re-raises it.
@@ -210,7 +201,7 @@ fn drain(pool: &Pool) {
             if observed >= n {
                 break;
             }
-            let batch = ((n - observed) / (NUM_THREADS * 2)).clamp(1, MAX_CLAIM_BATCH);
+            let batch = ((n - observed) / (nt * 2)).clamp(1, MAX_CLAIM_BATCH);
             let start = pool.counter.fetch_add(batch, Ordering::Relaxed);
             if start >= n {
                 break;
@@ -232,7 +223,8 @@ pub fn for_each_chunk<F: Fn(usize, usize) + Sync>(n_tasks: usize, f: F) {
     assert!(!IN_TASK.get(), "nested parallel dispatch from within a pool task");
 
     // Trivial sizes / single-core builds run inline.
-    if NUM_THREADS <= 1 || n_tasks <= 1 {
+    let nt = num_threads();
+    if nt <= 1 || n_tasks <= 1 {
         if n_tasks > 0 {
             f(0, n_tasks);
         }
@@ -252,7 +244,7 @@ pub fn for_each_chunk<F: Fn(usize, usize) + Sync>(n_tasks: usize, f: F) {
     // SAFETY: sole writer — prior dispatch fully drained (`working == 0`), next not yet observed.
     unsafe { *pool.job.get() = Some(Job { f: f_erased, n_tasks }) };
     pool.counter.store(0, Ordering::Relaxed);
-    pool.working.store(NUM_THREADS - 1, Ordering::Release);
+    pool.working.store(nt - 1, Ordering::Release);
     pool.generation.fetch_add(1, Ordering::SeqCst); // publish; SeqCst guards the park protocol
 
     // Wake only parked workers; spinning ones see the bump for free.
@@ -389,7 +381,7 @@ pub fn par_fill<T: Send, F: Fn(usize) -> T + Sync>(dst: &mut [T], build: F) {
 /// `run(slot, start, end)` fires once per claimed batch with that worker's slot, so state
 /// accumulates across its batches. Returns the slots (rest `None`) for the caller to combine.
 fn drain_into_slots<S: Send>(n_tasks: usize, run: impl Fn(&mut Option<S>, usize, usize) + Sync) -> Vec<Option<S>> {
-    let mut slots: Vec<Option<S>> = (0..NUM_THREADS).map(|_| None).collect();
+    let mut slots: Vec<Option<S>> = (0..num_threads()).map(|_| None).collect();
     let ptr = SendPtr(slots.as_mut_ptr());
     for_each_chunk(n_tasks, |start, end| {
         // SAFETY: `current_worker_id() < NUM_THREADS` is unique per live worker → disjoint

crates/backend/poly/src/eq_mle.rs

@@ -2,21 +2,20 @@ use crate::*;
 use crate::{EFPacking, PF};
 use ::utils::{iter_array_chunks_padded, log2_ceil_usize, log2_strict_usize};
 use field::*;
-use system_info::NUM_THREADS;
+use system_info::num_threads;
 use zk_alloc::ArenaVec;
 
-const LOG_NUM_THREADS: usize = log2_ceil_usize(NUM_THREADS);
 const LOG_BATCHED_TILE_SIZE: usize = 14;
 
-/// log2 oversubscription for the eq_mle fan-out: emit `NUM_THREADS << this` chunks so the
+/// log2 oversubscription for the eq_mle fan-out: emit `num_threads() << this` chunks so the
 /// pool's task counter rebalances across heterogeneous cores (e.g. P/E). `0` = one chunk
 /// per worker; `2` (4x) is a conservative default that balances well without over-fragmenting.
 const PARALLEL_LOG_OVERSUB: usize = 2;
 
 /// `(log2(n_chunks), n_chunks)` for the parallel fan-out.
 #[inline]
 fn parallel_split() -> (usize, usize) {
-    let log_chunks = LOG_NUM_THREADS + PARALLEL_LOG_OVERSUB;
+    let log_chunks = log2_ceil_usize(num_threads()) + PARALLEL_LOG_OVERSUB;
     (log_chunks, 1 << log_chunks)
 }
 

crates/backend/system-info/build.rs

@@ -1,11 +1,61 @@
-include!(concat!(env!("OUT_DIR"), "/info.rs"));
+use std::sync::OnceLock;
 
 const _: () = assert!(usize::BITS == 64, "this project requires a 64-bit target (for now)");
 
+#[must_use]
+pub fn num_threads() -> usize {
+    static CACHE: OnceLock<usize> = OnceLock::new();
+    *CACHE.get_or_init(|| {
+        std::thread::available_parallelism()
+            .expect("failed to detect available parallelism")
+            .get()
+    })
+}
+
+#[must_use]
+pub fn l1_cache_size() -> usize {
+    static CACHE: OnceLock<usize> = OnceLock::new();
+    *CACHE.get_or_init(|| {
+        detect_l1_cache_size().unwrap_or_else(|| {
+            eprintln!("Warning: failed to detect L1 cache size, defaulting to 32 KB");
+            32 * 1024
+        })
+    })
+}
+
 pub fn peak_rss_bytes() -> u64 {
     let mut ru: libc::rusage = unsafe { std::mem::zeroed() };
     unsafe { libc::getrusage(libc::RUSAGE_SELF, &raw mut ru) };
     let max = ru.ru_maxrss as u64;
     // ru_maxrss unit: bytes on macOS, KiB on Linux.
     if cfg!(target_os = "macos") { max } else { max * 1024 }
 }
+
+#[cfg(target_os = "linux")]
+fn detect_l1_cache_size() -> Option<usize> {
+    // /sys reports e.g. "32K\n", "48K\n", "1M\n".
+    let s = std::fs::read_to_string("/sys/devices/system/cpu/cpu0/cache/index0/size").ok()?;
+    let s = s.trim();
+    let last = s.chars().last()?;
+    match last {
+        'K' | 'k' => s[..s.len() - 1].parse::<usize>().ok().map(|n| n * 1024),
+        'M' | 'm' => s[..s.len() - 1].parse::<usize>().ok().map(|n| n * 1024 * 1024),
+        c if c.is_ascii_digit() => s.parse().ok(),
+        _ => None,
+    }
+}
+
+#[cfg(target_os = "macos")]
+fn detect_l1_cache_size() -> Option<usize> {
+    // `hw.l1dcachesize` returns the E-core value on Apple Silicon; prefer the P-core size.
+    let read_sysctl = |key: &str| -> Option<usize> {
+        let out = std::process::Command::new("sysctl").args(["-n", key]).output().ok()?;
+        std::str::from_utf8(&out.stdout).ok()?.trim().parse().ok()
+    };
+    read_sysctl("hw.perflevel0.l1dcachesize").or_else(|| read_sysctl("hw.l1dcachesize"))
+}
+
+#[cfg(not(any(target_os = "linux", target_os = "macos")))]
+fn detect_l1_cache_size() -> Option<usize> {
+    None
+}

crates/backend/system-info/src/lib.rs

@@ -8,7 +8,7 @@ use std::cell::Cell;
 use std::sync::OnceLock;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 
-use system_info::NUM_THREADS;
+use system_info::num_threads;
 
 mod arena_cow;
 mod arena_vec;
@@ -27,8 +27,15 @@ macro_rules! arena_vec {
 
 const SLAB_SIZE: usize = 8 << 30; // 8 GiB; per-thread soft cap, overflow falls back to System
 const SLACK: usize = 4; // extra slabs for non-pool threads that allocate in a phase
-const MAX_THREADS: usize = NUM_THREADS + SLACK;
-const REGION_SIZE: usize = SLAB_SIZE * MAX_THREADS; // one contiguous region => O(1) pointer classification
+
+fn max_threads() -> usize {
+    num_threads() + SLACK
+}
+
+fn region_size() -> usize {
+    static SIZE: OnceLock<usize> = OnceLock::new();
+    *SIZE.get_or_init(|| SLAB_SIZE * max_threads())
+}
 
 /// Bumped by `begin_phase()`; a thread resets its slab when its cached `ARENA_GEN` lags — one store
 /// resets every thread, lock-free.
@@ -58,12 +65,13 @@ thread_local! {
 
 fn ensure_region() -> usize {
     *REGION.get_or_init(|| {
+        let size = region_size();
         // SAFETY: mmap returns a page-aligned pointer or null; lazily backed.
-        let ptr = unsafe { syscall::mmap_anonymous(REGION_SIZE) };
+        let ptr = unsafe { syscall::mmap_anonymous(size) };
         if ptr.is_null() {
             std::process::abort();
         }
-        unsafe { syscall::madvise(ptr, REGION_SIZE, syscall::MADV_NOHUGEPAGE) };
+        unsafe { syscall::madvise(ptr, size, syscall::MADV_NOHUGEPAGE) };
         ptr as usize
     })
 }
@@ -125,7 +133,7 @@ unsafe fn arena_alloc_cold(size: usize, align: usize) -> *mut u8 {
         if base == 0 {
             let region = ensure_region();
             let idx = THREAD_IDX.fetch_add(1, Ordering::Relaxed);
-            if idx >= MAX_THREADS {
+            if idx >= max_threads() {
                 ARENA_NO_SLAB.set(true);
                 return unsafe { std::alloc::System.alloc(Layout::from_size_align_unchecked(size, align)) };
             }
@@ -177,7 +185,7 @@ pub(crate) unsafe fn raw_dealloc(ptr: *mut u8, size: usize, align: usize) {
     let addr = ptr as usize;
     if REGION
         .get()
-        .is_some_and(|&base| addr >= base && addr < base + REGION_SIZE)
+        .is_some_and(|&base| addr >= base && addr < base + region_size())
     {
         return; // arena pointer — free is a no-op
     }

crates/backend/zk-alloc/src/lib.rs

@@ -419,8 +419,8 @@ fn fft_triple_layer_quad_twiddle<F: Field, Fly: Butterfly<F>>(
 
 /// Estimates the optimal workload size for `T` to fit in L1 cache.
 #[must_use]
-const fn workload_size<T: Sized>() -> usize {
-    system_info::L1_CACHE_SIZE / size_of::<T>()
+fn workload_size<T: Sized>() -> usize {
+    system_info::l1_cache_size() / size_of::<T>()
 }
 
 /// Estimates the optimal number of rows of a `RowMajorMatrix<T>` to take in each parallel chunk.

crates/whir/src/dft.rs

@@ -183,7 +183,7 @@ fn prepare_evals_for_fft_unpacked<A: Copy + Send + Sync>(
         return out;
     }
 
-    let rows_per_band = ((system_info::L1_CACHE_SIZE / 2) / (dft_n_cols * size_of::<A>())).clamp(1, block_size);
+    let rows_per_band = ((system_info::l1_cache_size() / 2) / (dft_n_cols * size_of::<A>())).clamp(1, block_size);
     let band_len = rows_per_band * dft_n_cols;
 
     parallel::par_chunks_mut(&mut out, band_len, |band_idx, band| {
@@ -220,7 +220,7 @@ fn prepare_evals_for_fft_packed_extension<EF: ExtensionField<PF<EF>>>(
         return out;
     }
 
-    let rows_per_band = ((system_info::L1_CACHE_SIZE / 2) / (n_blocks * size_of::<EF>())).clamp(1, block_size);
+    let rows_per_band = ((system_info::l1_cache_size() / 2) / (n_blocks * size_of::<EF>())).clamp(1, block_size);
     let band_len = rows_per_band * n_blocks;
 
     parallel::par_chunks_mut(&mut out, band_len, |band_idx, band| {