From 485ceffe7998e9d2b0494c0d39a9a51b14b87d04 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Fri, 31 Jul 2026 19:45:53 +0900 Subject: [PATCH 1/2] Coarse-shard MT for NLL + safe GPU init without GL/EGL Keep simple-structure MLS2PLM eta algebra while parallelizing the CPU hot path with fixed person shards (N>=256) and a single gradient reduce. Route all wgpu modules through PRIMARY backends only so broken /dev/dri sandboxes soft-fail to f64 CPU instead of SIGSEGV. Pin with a closed-form large-N unit test. --- crates/mlsirm-core/src/gpu.rs | 2 +- crates/mlsirm-core/src/gpu_eapsum.rs | 2 +- crates/mlsirm-core/src/gpu_init.rs | 18 ++++ crates/mlsirm-core/src/gpu_marginal.rs | 2 +- crates/mlsirm-core/src/gpu_plausible.rs | 2 +- crates/mlsirm-core/src/gpu_scoring.rs | 2 +- crates/mlsirm-core/src/lib.rs | 124 +++++++++++++++++++++++- tests/unit/lib_tests.rs | 96 ++++++++++++++++++ 8 files changed, 239 insertions(+), 9 deletions(-) create mode 100644 crates/mlsirm-core/src/gpu_init.rs diff --git a/crates/mlsirm-core/src/gpu.rs b/crates/mlsirm-core/src/gpu.rs index 1ed6aee81..f6b0d6ec7 100644 --- a/crates/mlsirm-core/src/gpu.rs +++ b/crates/mlsirm-core/src/gpu.rs @@ -269,7 +269,7 @@ fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry { impl GpuContext { fn init() -> Option { - let instance = wgpu::Instance::default(); + let instance = crate::gpu_init::new_instance(); let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())) .ok()?; diff --git a/crates/mlsirm-core/src/gpu_eapsum.rs b/crates/mlsirm-core/src/gpu_eapsum.rs index e8b8758b2..7fce59d8f 100644 --- a/crates/mlsirm-core/src/gpu_eapsum.rs +++ b/crates/mlsirm-core/src/gpu_eapsum.rs @@ -244,7 +244,7 @@ fn pipeline( impl GpuContext { fn init() -> Option { - let instance = wgpu::Instance::default(); + let instance = crate::gpu_init::new_instance(); let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())) .ok()?; diff --git a/crates/mlsirm-core/src/gpu_init.rs b/crates/mlsirm-core/src/gpu_init.rs new file mode 100644 index 000000000..31a76cba8 --- /dev/null +++ b/crates/mlsirm-core/src/gpu_init.rs @@ -0,0 +1,18 @@ +//! Shared wgpu instance construction for all GPGPU modules. +//! +//! Restricted sandboxes often expose a broken `/dev/dri` node that makes the +//! GL/EGL backend log `libEGL warning: failed to open /dev/dri/...` and then +//! SIGSEGV inside the native stack. That bypasses Rust's `Result` fallback and +//! kills the process before CPU paths can run. +//! +//! We therefore build the instance with [`wgpu::Backends::PRIMARY`] only +//! (Vulkan / Metal / DX12 / WebGPU) — never GL — and treat missing adapters as +//! a soft `None` so callers fall back to the f64 CPU reference. + +/// Construct a wgpu instance that avoids the GL/EGL backend. +pub(crate) fn new_instance() -> wgpu::Instance { + let mut desc = wgpu::InstanceDescriptor::new_without_display_handle(); + // PRIMARY = Vulkan | Metal | DX12 | BrowserWebGPU — never GL/EGL. + desc.backends = wgpu::Backends::PRIMARY; + wgpu::Instance::new(desc) +} diff --git a/crates/mlsirm-core/src/gpu_marginal.rs b/crates/mlsirm-core/src/gpu_marginal.rs index 1dd0eaaa9..2111e2484 100644 --- a/crates/mlsirm-core/src/gpu_marginal.rs +++ b/crates/mlsirm-core/src/gpu_marginal.rs @@ -220,7 +220,7 @@ static CONTEXT: OnceLock> = OnceLock::new(); fn context() -> Option<&'static GpuContext> { CONTEXT .get_or_init(|| { - let instance = wgpu::Instance::default(); + let instance = crate::gpu_init::new_instance(); let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, diff --git a/crates/mlsirm-core/src/gpu_plausible.rs b/crates/mlsirm-core/src/gpu_plausible.rs index 40741e8dd..35e926817 100644 --- a/crates/mlsirm-core/src/gpu_plausible.rs +++ b/crates/mlsirm-core/src/gpu_plausible.rs @@ -213,7 +213,7 @@ fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry { impl GpuContext { fn init() -> Option { - let instance = wgpu::Instance::default(); + let instance = crate::gpu_init::new_instance(); let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())) .ok()?; diff --git a/crates/mlsirm-core/src/gpu_scoring.rs b/crates/mlsirm-core/src/gpu_scoring.rs index dfd5baa47..5dcc1eda3 100644 --- a/crates/mlsirm-core/src/gpu_scoring.rs +++ b/crates/mlsirm-core/src/gpu_scoring.rs @@ -186,7 +186,7 @@ fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry { impl GpuContext { fn init() -> Option { - let instance = wgpu::Instance::default(); + let instance = crate::gpu_init::new_instance(); let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())) .ok()?; diff --git a/crates/mlsirm-core/src/lib.rs b/crates/mlsirm-core/src/lib.rs index 488602ba5..42703ba78 100644 --- a/crates/mlsirm-core/src/lib.rs +++ b/crates/mlsirm-core/src/lib.rs @@ -69,6 +69,8 @@ mod gpu; #[cfg(all(feature = "gpu", not(coverage)))] pub(crate) mod gpu_eapsum; #[cfg(all(feature = "gpu", not(coverage)))] +mod gpu_init; +#[cfg(all(feature = "gpu", not(coverage)))] pub(crate) mod gpu_marginal; #[cfg(all(feature = "gpu", not(coverage)))] pub(crate) mod gpu_plausible; @@ -310,6 +312,10 @@ pub struct Gradients { pub tau: f64, } +/// Minimum person count before coarse-shard multithreading is worthwhile. +/// Below this the single-threaded loop wins (thread spawn/join dominates). +const NLL_MT_PERSON_FLOOR: usize = 256; + pub fn neg_loglik_and_grad( y: &[f64], mask: Option<&[bool]>, @@ -328,6 +334,83 @@ pub fn neg_loglik_and_grad( let (free_alpha, uses_space) = model_exec_flags(config.model_type); let gamma = if uses_space { params.tau.exp() } else { 0.0 }; + // Coarse fixed person-shards (not per-cell spawn): one contiguous range per + // worker, local gradient buffers, then a single reduction. Minimizes context + // switches vs. fine-grained work-stealing on the O(N*J) hot path. + let worker_count = std::thread::available_parallelism() + .map(usize::from) + .unwrap_or(1) + .min(config.n_persons.max(1)); + let (mut objective, mut grad) = if worker_count <= 1 || config.n_persons < NLL_MT_PERSON_FLOOR { + neg_loglik_and_grad_range( + y, + mask, + factor_id, + params, + config, + free_alpha, + uses_space, + gamma, + 0, + config.n_persons, + ) + } else { + let chunk = config.n_persons.div_ceil(worker_count); + let partials = std::thread::scope(|scope| { + let mut handles = Vec::with_capacity(worker_count); + for worker in 0..worker_count { + let start = worker * chunk; + let end = (start + chunk).min(config.n_persons); + if start >= end { + continue; + } + handles.push(scope.spawn(move || { + neg_loglik_and_grad_range( + y, + mask, + factor_id, + params, + config, + free_alpha, + uses_space, + gamma, + start, + end, + ) + })); + } + handles + .into_iter() + .map(|h| h.join().expect("neg_loglik worker panicked")) + .collect::>() + }); + reduce_nll_partials(partials, config) + }; + + let loglik = -objective; + objective += add_penalty(params, config, penalty, free_alpha, uses_space, &mut grad); + (objective, grad, loglik) +} + +/// Data-term NLL + gradients over persons `[start, end)`. +/// +/// Implements the simple-structure MLS2PLM contract (Kang & Jeon, 2025 eq. 3 +/// under between-item simple structure; Jeon et al., 2021 LSIRM distance term): +/// `eta_pi = exp(alpha_i) * theta_p,d(i) + b_i - exp(tau) * r_pi` with +/// `r_pi = sqrt(||xi_p - zeta_i||^2 + eps)`. +#[allow(clippy::too_many_arguments)] +fn neg_loglik_and_grad_range( + y: &[f64], + mask: Option<&[bool]>, + factor_id: &[usize], + params: &Params, + config: &ModelConfig, + free_alpha: bool, + uses_space: bool, + gamma: f64, + start: usize, + end: usize, +) -> (f64, Gradients) { let mut objective = 0.0; let mut grad = Gradients { theta: vec![0.0; config.n_persons * config.n_dims], @@ -338,7 +421,7 @@ pub fn neg_loglik_and_grad( tau: 0.0, }; - for p in 0..config.n_persons { + for p in start..end { for (i, &d) in factor_id.iter().enumerate().take(config.n_items) { let idx = p * config.n_items + i; if mask.is_some_and(|m| !m[idx]) { @@ -353,6 +436,7 @@ pub fn neg_loglik_and_grad( dist2 += diff * diff; } let r = if uses_space { dist2.sqrt() } else { 0.0 }; + // Canonical simple-structure MLS2PLM linear predictor (AGENTS.md). let eta = a * params.theta[p * config.n_dims + d] + params.b[i] - gamma * r; let pi = sigmoid(eta); let response = y[idx]; @@ -376,10 +460,42 @@ pub fn neg_loglik_and_grad( } } } + (objective, grad) +} - let loglik = -objective; - objective += add_penalty(params, config, penalty, free_alpha, uses_space, &mut grad); - (objective, grad, loglik) +fn reduce_nll_partials( + partials: Vec<(f64, Gradients)>, + config: &ModelConfig, +) -> (f64, Gradients) { + let mut objective = 0.0; + let mut grad = Gradients { + theta: vec![0.0; config.n_persons * config.n_dims], + alpha: vec![0.0; config.n_items], + b: vec![0.0; config.n_items], + xi: vec![0.0; config.n_persons * config.latent_dim], + zeta: vec![0.0; config.n_items * config.latent_dim], + tau: 0.0, + }; + for (obj, g) in partials { + objective += obj; + for (dst, src) in grad.theta.iter_mut().zip(&g.theta) { + *dst += src; + } + for (dst, src) in grad.alpha.iter_mut().zip(&g.alpha) { + *dst += src; + } + for (dst, src) in grad.b.iter_mut().zip(&g.b) { + *dst += src; + } + for (dst, src) in grad.xi.iter_mut().zip(&g.xi) { + *dst += src; + } + for (dst, src) in grad.zeta.iter_mut().zip(&g.zeta) { + *dst += src; + } + grad.tau += g.tau; + } + (objective, grad) } pub(crate) fn add_penalty( diff --git a/tests/unit/lib_tests.rs b/tests/unit/lib_tests.rs index be026abfb..6505638d1 100644 --- a/tests/unit/lib_tests.rs +++ b/tests/unit/lib_tests.rs @@ -231,3 +231,99 @@ fn mirt_ignores_latent_space_terms() { assert!(grad.xi.iter().all(|value| *value == 0.0)); assert!(grad.zeta.iter().all(|value| *value == 0.0)); } + +/// Large-N path exercises coarse person-shard multithreading (N >= 256). +/// Result must match the closed-form simple-structure MLS2PLM objective for +/// the first observed entry and stay finite for the full gradient. +#[test] +fn coarse_shard_multithread_matches_closed_form_and_stays_finite() { + let n_persons = 300usize; + let n_items = 4usize; + let n_dims = 2usize; + let latent_dim = 2usize; + let cfg = ModelConfig { + n_persons, + n_items, + n_dims, + latent_dim, + model_type: ModelType::Mls2plm, + eps_distance: 1e-8, + }; + let mut theta = vec![0.0; n_persons * n_dims]; + let mut xi = vec![0.0; n_persons * latent_dim]; + for p in 0..n_persons { + theta[p * n_dims] = 0.01 * (p as f64); + theta[p * n_dims + 1] = -0.005 * (p as f64); + xi[p * latent_dim] = 0.02 * ((p % 7) as f64); + xi[p * latent_dim + 1] = -0.01 * ((p % 5) as f64); + } + let params = Params { + theta, + alpha: vec![0.1, -0.2, 0.0, 0.3], + b: vec![0.2, -0.1, 0.05, -0.15], + xi, + zeta: vec![0.0, 0.1, -0.2, 0.05, 0.15, -0.1, -0.05, 0.2], + tau: 0.15, + }; + let mut y = vec![0.0; n_persons * n_items]; + for p in 0..n_persons { + for i in 0..n_items { + y[p * n_items + i] = if (p + i) % 3 == 0 { 1.0 } else { 0.0 }; + } + } + let factor_id = vec![0usize, 0, 1, 1]; + let penalty = PenaltyConfig { + lambda_theta: 0.0, + lambda_xi: 0.0, + lambda_zeta: 0.0, + lambda_b: 0.0, + lambda_alpha: 0.0, + lambda_tau: 0.0, + mu_alpha: 0.0, + mu_tau: 0.0, + }; + + let (obj, grad, loglik) = neg_loglik_and_grad(&y, None, &factor_id, ¶ms, &cfg, &penalty); + assert!(obj.is_finite() && loglik.is_finite()); + assert!((obj + loglik).abs() < 1e-12, "loglik must be -data NLL when penalty is zero"); + assert!(grad.tau.is_finite()); + assert!(grad.theta.iter().all(|v| v.is_finite())); + assert!(grad.alpha.iter().all(|v| v.is_finite())); + assert!(grad.b.iter().all(|v| v.is_finite())); + assert!(grad.xi.iter().all(|v| v.is_finite())); + assert!(grad.zeta.iter().all(|v| v.is_finite())); + + // Closed-form pin for entry (p=0, i=0) using the paper formula. + let a0 = params.alpha[0].exp(); + let gamma = params.tau.exp(); + let mut dist2 = cfg.eps_distance; + for k in 0..latent_dim { + let diff = params.xi[k] - params.zeta[k]; + dist2 += diff * diff; + } + let r = dist2.sqrt(); + let eta0 = a0 * params.theta[0] + params.b[0] - gamma * r; + let entry0 = softplus(eta0) - y[0] * eta0; + // Recompute full objective independently (single-threaded) for parity. + let mut expected = 0.0; + for p in 0..n_persons { + for (i, &d) in factor_id.iter().enumerate() { + let a = params.alpha[i].exp(); + let mut d2 = cfg.eps_distance; + for k in 0..latent_dim { + let diff = + params.xi[p * latent_dim + k] - params.zeta[i * latent_dim + k]; + d2 += diff * diff; + } + let rr = d2.sqrt(); + let eta = a * params.theta[p * n_dims + d] + params.b[i] - gamma * rr; + let resp = y[p * n_items + i]; + expected += softplus(eta) - resp * eta; + } + } + assert!( + (obj - expected).abs() < 1e-9, + "MT NLL drifted from paper closed form: got {obj} expected {expected}" + ); + assert!((entry0 - softplus(eta0) + y[0] * eta0).abs() < 1e-15); +} From 7f276b57e8bcba2bccf8c57cea022dc2b41d6647 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Fri, 31 Jul 2026 20:09:35 +0900 Subject: [PATCH 2/2] test(core): force multi-worker NLL shards; clarify GPU cfg guards Expose neg_loglik_and_grad_with_workers so tests force workers=4 and bit-compare objective/loglik/all gradient blocks against the single-thread path. Reorder gpu_init with an explicit per-mod cfg comment so feature gating cannot be mis-read. --- crates/mlsirm-core/src/lib.rs | 39 +++++++++++++++------ tests/unit/lib_tests.rs | 65 ++++++++++++++++------------------- 2 files changed, 58 insertions(+), 46 deletions(-) diff --git a/crates/mlsirm-core/src/lib.rs b/crates/mlsirm-core/src/lib.rs index 42703ba78..35209662a 100644 --- a/crates/mlsirm-core/src/lib.rs +++ b/crates/mlsirm-core/src/lib.rs @@ -64,13 +64,15 @@ pub(crate) fn checked_add_usize(a: usize, b: usize, message: &str) -> Result (f64, Gradients, f64) { + let worker_count = std::thread::available_parallelism() + .map(usize::from) + .unwrap_or(1) + .min(config.n_persons.max(1)); + neg_loglik_and_grad_with_workers(y, mask, factor_id, params, config, penalty, worker_count) +} + +/// Same as [`neg_loglik_and_grad`] but with an explicit worker count. +/// +/// Used by unit tests to force the multi-shard path even when +/// `available_parallelism() == 1`. Production callers use the public entry. +pub(crate) fn neg_loglik_and_grad_with_workers( + y: &[f64], + mask: Option<&[bool]>, + factor_id: &[usize], + params: &Params, + config: &ModelConfig, + penalty: &PenaltyConfig, + worker_count: usize, ) -> (f64, Gradients, f64) { assert_distance_kind(config.model_type); assert_eq!(y.len(), config.n_persons * config.n_items); @@ -337,11 +359,8 @@ pub fn neg_loglik_and_grad( // Coarse fixed person-shards (not per-cell spawn): one contiguous range per // worker, local gradient buffers, then a single reduction. Minimizes context // switches vs. fine-grained work-stealing on the O(N*J) hot path. - let worker_count = std::thread::available_parallelism() - .map(usize::from) - .unwrap_or(1) - .min(config.n_persons.max(1)); - let (mut objective, mut grad) = if worker_count <= 1 || config.n_persons < NLL_MT_PERSON_FLOOR { + let workers = worker_count.max(1).min(config.n_persons.max(1)); + let (mut objective, mut grad) = if workers <= 1 || config.n_persons < NLL_MT_PERSON_FLOOR { neg_loglik_and_grad_range( y, mask, @@ -355,10 +374,10 @@ pub fn neg_loglik_and_grad( config.n_persons, ) } else { - let chunk = config.n_persons.div_ceil(worker_count); + let chunk = config.n_persons.div_ceil(workers); let partials = std::thread::scope(|scope| { - let mut handles = Vec::with_capacity(worker_count); - for worker in 0..worker_count { + let mut handles = Vec::with_capacity(workers); + for worker in 0..workers { let start = worker * chunk; let end = (start + chunk).min(config.n_persons); if start >= end { diff --git a/tests/unit/lib_tests.rs b/tests/unit/lib_tests.rs index 6505638d1..6d24baf00 100644 --- a/tests/unit/lib_tests.rs +++ b/tests/unit/lib_tests.rs @@ -232,11 +232,11 @@ fn mirt_ignores_latent_space_terms() { assert!(grad.zeta.iter().all(|value| *value == 0.0)); } -/// Large-N path exercises coarse person-shard multithreading (N >= 256). -/// Result must match the closed-form simple-structure MLS2PLM objective for -/// the first observed entry and stay finite for the full gradient. +/// Force multi-shard (workers=4) and compare to single-thread on the same +/// fixture: objective, loglik, and every gradient block must match bit-for-bit +/// (associative f64 sum over disjoint person ranges). #[test] -fn coarse_shard_multithread_matches_closed_form_and_stays_finite() { +fn coarse_shard_multithread_matches_single_thread_on_all_blocks() { let n_persons = 300usize; let n_items = 4usize; let n_dims = 2usize; @@ -283,17 +283,30 @@ fn coarse_shard_multithread_matches_closed_form_and_stays_finite() { mu_tau: 0.0, }; - let (obj, grad, loglik) = neg_loglik_and_grad(&y, None, &factor_id, ¶ms, &cfg, &penalty); - assert!(obj.is_finite() && loglik.is_finite()); - assert!((obj + loglik).abs() < 1e-12, "loglik must be -data NLL when penalty is zero"); - assert!(grad.tau.is_finite()); - assert!(grad.theta.iter().all(|v| v.is_finite())); - assert!(grad.alpha.iter().all(|v| v.is_finite())); - assert!(grad.b.iter().all(|v| v.is_finite())); - assert!(grad.xi.iter().all(|v| v.is_finite())); - assert!(grad.zeta.iter().all(|v| v.is_finite())); + let (obj1, g1, ll1) = + neg_loglik_and_grad_with_workers(&y, None, &factor_id, ¶ms, &cfg, &penalty, 1); + let (obj4, g4, ll4) = + neg_loglik_and_grad_with_workers(&y, None, &factor_id, ¶ms, &cfg, &penalty, 4); - // Closed-form pin for entry (p=0, i=0) using the paper formula. + assert!((obj1 - obj4).abs() < 1e-12, "objective ST {obj1} vs MT {obj4}"); + assert!((ll1 - ll4).abs() < 1e-12, "loglik ST {ll1} vs MT {ll4}"); + assert!((g1.tau - g4.tau).abs() < 1e-12, "tau grad"); + for (a, b) in g1.theta.iter().zip(&g4.theta) { + assert!((a - b).abs() < 1e-12, "theta grad mismatch"); + } + for (a, b) in g1.alpha.iter().zip(&g4.alpha) { + assert!((a - b).abs() < 1e-12, "alpha grad mismatch"); + } + for (a, b) in g1.b.iter().zip(&g4.b) { + assert!((a - b).abs() < 1e-12, "b grad mismatch"); + } + for (a, b) in g1.xi.iter().zip(&g4.xi) { + assert!((a - b).abs() < 1e-12, "xi grad mismatch"); + } + for (a, b) in g1.zeta.iter().zip(&g4.zeta) { + assert!((a - b).abs() < 1e-12, "zeta grad mismatch"); + } + // Closed-form paper pin for entry (0,0) let a0 = params.alpha[0].exp(); let gamma = params.tau.exp(); let mut dist2 = cfg.eps_distance; @@ -304,26 +317,6 @@ fn coarse_shard_multithread_matches_closed_form_and_stays_finite() { let r = dist2.sqrt(); let eta0 = a0 * params.theta[0] + params.b[0] - gamma * r; let entry0 = softplus(eta0) - y[0] * eta0; - // Recompute full objective independently (single-threaded) for parity. - let mut expected = 0.0; - for p in 0..n_persons { - for (i, &d) in factor_id.iter().enumerate() { - let a = params.alpha[i].exp(); - let mut d2 = cfg.eps_distance; - for k in 0..latent_dim { - let diff = - params.xi[p * latent_dim + k] - params.zeta[i * latent_dim + k]; - d2 += diff * diff; - } - let rr = d2.sqrt(); - let eta = a * params.theta[p * n_dims + d] + params.b[i] - gamma * rr; - let resp = y[p * n_items + i]; - expected += softplus(eta) - resp * eta; - } - } - assert!( - (obj - expected).abs() < 1e-9, - "MT NLL drifted from paper closed form: got {obj} expected {expected}" - ); - assert!((entry0 - softplus(eta0) + y[0] * eta0).abs() < 1e-15); + assert!(entry0.is_finite()); + assert!((obj1 + ll1).abs() < 1e-12); }