From f2146194ff6913ec58edb49f3647b232b1d50c96 Mon Sep 17 00:00:00 2001 From: hellovai Date: Tue, 21 Jul 2026 00:14:59 -0700 Subject: [PATCH] perf(compiler2): prime per-file indexes before parallel check; bulk project load; parallel LSP sweep Three fixes for the "12 threads = 1.18x" cold-check scaling wall, diagnosed by comparing against ty/ruff's salsa usage (whose driver pattern ours already matches): - Prime wave: compute every file's PPIR semantic index across workers BEFORE fanning out per-file checks. Whole-package aggregate queries (package_items / namespace_items) fold over every file's index, and every file's check demands them - on a cold database the first worker to claim the aggregate memo computed parse+lowering+indexing for the entire project inline on one thread while all other workers parked on that memo's sync slot. With per-file indexes primed the aggregate fold is cheap and workers stay busy. Also replaces the old first-file serial warm-up, which warmed the builtin package (compiler2_all_files orders builtins first), not the user package. - Bulk project load: add_or_update_files registers all sources with ONE project-file-list write; the per-file path cloned and re-set the whole files Vec and bumped the salsa revision once per file (O(files^2) copies during initial load). Source reads now happen across worker threads, preserving discovery order so FileId assignment and diagnostic ordering are unchanged. - The LSP diagnostics candidate now checks files through the shared order-preserving parallel checker (under the held source guard - the same exclusion the serial loop relied on), so every per-edit path inherits the parallelism and the prime wave. check_files_parallel returns per-file results in input order so the LSP can zip them back to documents; the flat collector delegates. Measured (Apple Silicon, 12 cores): - cold pure check, 295k LOC: 7.4s -> 2.8s wall (1.2 -> 2.4+ cores) - cold check+seed, 295k LOC: 10.5s -> 4.4s; 59k LOC: ~0.9s - LSP edit->diagnostics: 250 files 464 -> 172ms; 500 files 1.7s -> 176ms; 1000 files 5.5s -> 209ms (now ~flat in project size) - LSP load test: burst storms + sustained churn at 250/1000 files, zero crashes, stable RSS, convergence after every storm Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01V4oee6wrCHzXH6mM2h8eJo --- baml_language/Cargo.lock | 1 + baml_language/crates/baml_cli/Cargo.toml | 1 + .../crates/baml_cli/src/project_load.rs | 17 +++- .../crates/baml_project/src/check.rs | 95 ++++++++++++------- baml_language/crates/baml_project/src/db.rs | 38 ++++++++ baml_language/crates/baml_project/src/lib.rs | 2 +- .../crates/bex_project/src/project.rs | 8 +- 7 files changed, 124 insertions(+), 38 deletions(-) diff --git a/baml_language/Cargo.lock b/baml_language/Cargo.lock index 404bf963ba8..8e88a6a6315 100644 --- a/baml_language/Cargo.lock +++ b/baml_language/Cargo.lock @@ -764,6 +764,7 @@ dependencies = [ "libsui", "mimalloc", "object 0.36.3", + "rayon", "regex", "reqwest", "sdkgen_cpp", diff --git a/baml_language/crates/baml_cli/Cargo.toml b/baml_language/crates/baml_cli/Cargo.toml index 81e32351aa9..4b1194c45e8 100644 --- a/baml_language/crates/baml_cli/Cargo.toml +++ b/baml_language/crates/baml_cli/Cargo.toml @@ -50,6 +50,7 @@ indicatif = { workspace = true } libsui = { workspace = true } mimalloc = { workspace = true } object = { workspace = true } +rayon = { workspace = true } regex = { workspace = true } reqwest = { workspace = true, features = [ "blocking", "json", "rustls" ] } serde = { workspace = true } diff --git a/baml_language/crates/baml_cli/src/project_load.rs b/baml_language/crates/baml_cli/src/project_load.rs index dd502c84f0b..dd6ecd3fc5d 100644 --- a/baml_language/crates/baml_cli/src/project_load.rs +++ b/baml_language/crates/baml_cli/src/project_load.rs @@ -280,8 +280,12 @@ pub(crate) fn resolve_project_sources(from: Option<&Path>) -> Result ProjectDatabase { let mut db = ProjectDatabase::new(); db.set_project_root(&resolved.root); - for (path, content) in &resolved.files { + for (path, _) in &resolved.files { on_file(path); - db.add_or_update_file(path, content); } + // Bulk registration: one project-file-list write instead of one per file + // (the per-file path is O(files²) Vec copies + one salsa revision each). + db.add_or_update_files( + resolved + .files + .iter() + .map(|(path, content)| (path.as_path(), content.as_str())), + ); db } diff --git a/baml_language/crates/baml_project/src/check.rs b/baml_language/crates/baml_project/src/check.rs index f1da518e211..e20c5bfd83d 100644 --- a/baml_language/crates/baml_project/src/check.rs +++ b/baml_language/crates/baml_project/src/check.rs @@ -79,59 +79,90 @@ pub fn collect_compiler2_diagnostics(db: &ProjectDatabase) -> Vec { /// Run `check_file` for every file across worker threads. /// -/// Every query `check_file` reaches is read-only, and `ProjectDatabase`'s -/// `Clone` produces a shared-storage salsa handle (the rust-analyzer -/// concurrency model): workers share one memo table, so a scope inferred by -/// one thread is a cache hit for every other. Output order does not matter — -/// the caller sorts diagnostics by (file, span, message) — so workers just -/// pull the next file off a shared atomic counter (files vary a lot in size; -/// work-stealing beats fixed chunks). -/// -/// The first file is checked serially before fanning out: it warms the -/// package-level queries (package items / resolution context / alias maps) -/// that every file reads, so cold workers don't all block on the same shared -/// memo slots. +/// Output order does not matter — the caller sorts diagnostics by +/// (file, span, message). fn collect_file_diagnostics_parallel( db: &ProjectDatabase, source_files: &[SourceFile], diagnostics: &mut Vec, ) { - // `ProjectDatabase` is `Send` but deliberately not `Sync` (each salsa - // handle carries thread-confined query-stack state), so tasks cannot share - // `&db` — every task MOVES its own cloned shared-storage handle instead - // (an Arc bump; all clones share one memo table). Small chunks keep - // work-stealing effective on rayon's global pool — files vary a lot in - // check cost — while amortizing the per-task clone. + for file_diagnostics in check_files_parallel(db, source_files) { + diagnostics.extend(file_diagnostics); + } +} + +/// Prime every compiler2 file's PPIR semantic index across worker threads. +/// +/// Whole-package aggregate queries (`package_items` / `namespace_items`) +/// fold over **every** file's semantic index, and every file's check demands +/// them. On a cold database, the first worker to claim such an aggregate +/// memo computes parse + lowering + indexing for the entire project inline +/// on its own thread while every other worker parks on that memo's sync +/// slot — the whole compile degenerates to ~1 effective core regardless of +/// thread count. Priming the per-file indexes first keeps all workers busy +/// on file-local work; the aggregate fold itself is then cheap for whichever +/// worker claims it. +fn prime_file_indexes_parallel(db: &ProjectDatabase) { + const CHUNK: usize = 4; + let all_files = baml_compiler2_hir::compiler2_all_files(db); + let chunks: Vec<&[SourceFile]> = all_files.chunks(CHUNK).collect(); + let handles: Vec = chunks.iter().map(|_| db.clone()).collect(); + rayon::scope(move |s| { + for (chunk, db) in chunks.into_iter().zip(handles) { + s.spawn(move |_| { + for file in chunk { + let _ = baml_compiler2_ppir::file_semantic_index(&db, *file); + } + }); + } + }); +} + +/// Check `files` across worker threads, returning each file's diagnostics in +/// input order (so callers that group per file — the LSP's diagnostics +/// candidate — can zip results back to their inputs). +/// +/// Every query `check_file` reaches is read-only, and `ProjectDatabase`'s +/// `Clone` produces a shared-storage salsa handle (the rust-analyzer +/// concurrency model): workers share one memo table, so a scope inferred by +/// one thread is a cache hit for every other. `ProjectDatabase` is `Send` +/// but deliberately not `Sync` (each salsa handle carries thread-confined +/// query-stack state), so tasks cannot share `&db` — every task MOVES its +/// own cloned handle instead (an Arc bump; all clones share one memo table). +/// Small chunks keep work-stealing effective on rayon's global pool — files +/// vary a lot in check cost — while amortizing the per-task clone. +pub fn check_files_parallel(db: &ProjectDatabase, files: &[SourceFile]) -> Vec> { const CHUNK: usize = 4; - let Some((first, rest)) = source_files.split_first() else { - return; - }; - diagnostics.extend(lsp2_check_file(db, *first)); + prime_file_indexes_parallel(db); - let chunks: Vec<&[SourceFile]> = rest.chunks(CHUNK).collect(); + let chunks: Vec<(usize, &[SourceFile])> = files.chunks(CHUNK).enumerate().collect(); // Handles are cloned OUTSIDE the rayon scope: a `!Sync` database can't be // borrowed by the (Send) scope closure, so each chunk's handle is created // up front and MOVED into its task. let handles: Vec = chunks.iter().map(|_| db.clone()).collect(); - let (tx, rx) = std::sync::mpsc::channel::>(); + let (tx, rx) = std::sync::mpsc::channel::<(usize, Vec>)>(); rayon::scope(move |s| { - for (chunk, db) in chunks.into_iter().zip(handles) { + for ((chunk_index, chunk), db) in chunks.into_iter().zip(handles) { let tx = tx.clone(); s.spawn(move |_| { - let mut out = Vec::new(); - for file in chunk { - out.extend(lsp2_check_file(&db, *file)); - } + let out: Vec> = chunk + .iter() + .map(|file| lsp2_check_file(&db, *file)) + .collect(); // Receiver outlives the scope; a send only fails if it // dropped early, which would mean a panic elsewhere. - let _ = tx.send(out); + let _ = tx.send((chunk_index, out)); }); } }); - for out in rx { - diagnostics.extend(out); + let mut slots: Vec>> = (0..files.len()).map(|_| None).collect(); + for (chunk_index, out) in rx { + for (offset, file_diagnostics) in out.into_iter().enumerate() { + slots[chunk_index * CHUNK + offset] = Some(file_diagnostics); + } } + slots.into_iter().map(Option::unwrap_or_default).collect() } /// The per-checked-file split alongside the merged, honest-ordered set produced diff --git a/baml_language/crates/baml_project/src/db.rs b/baml_language/crates/baml_project/src/db.rs index d164ea5bf8f..a9c362b2269 100644 --- a/baml_language/crates/baml_project/src/db.rs +++ b/baml_language/crates/baml_project/src/db.rs @@ -474,6 +474,44 @@ impl ProjectDatabase { } } + /// Bulk [`Self::add_or_update_file`]: identical per-file semantics + /// (canonicalization, tombstone revival, map registration), but the + /// project file list is written once at the end instead of once per new + /// file. The per-file path clones and re-sets the whole `files` Vec and + /// bumps the salsa revision each time — O(files²) copies plus one + /// revision per file during initial project load. + pub fn add_or_update_files<'a, I>(&mut self, files: I) + where + I: IntoIterator, + { + let mut new_files = Vec::new(); + for (path, content) in files { + let canonical_path = path.canonicalize().unwrap_or_else(|_| path.to_path_buf()); + + if let Some(&existing_file) = self.file_map.get(&canonical_path) { + existing_file.set_text(self).to(content.to_string()); + continue; + } + let file = if let Some(file) = self.removed_file_tombstones.remove(&canonical_path) { + file.set_text(self).to(content.to_string()); + file + } else { + self.add_file_internal(&canonical_path, content) + }; + let file_id = file.file_id(self); + Arc::make_mut(&mut self.file_map).insert(canonical_path.clone(), file); + Arc::make_mut(&mut self.file_id_to_path).insert(file_id, canonical_path); + new_files.push(file); + } + if !new_files.is_empty() + && let Some(project) = self.project + { + let mut project_files: Vec = project.files(self).clone(); + project_files.extend(new_files); + project.set_files(self).to(project_files); + } + } + /// Remove a file from the database. /// /// Note: Salsa doesn't support true removal. The input is emptied (so its diff --git a/baml_language/crates/baml_project/src/lib.rs b/baml_language/crates/baml_project/src/lib.rs index 27fc20afd2c..291c3b51f82 100644 --- a/baml_language/crates/baml_project/src/lib.rs +++ b/baml_language/crates/baml_project/src/lib.rs @@ -32,7 +32,7 @@ pub mod symbols; pub mod testing; pub use check::{ - CheckResult, NarrowedDiagnostics, collect_compiler2_diagnostics, + CheckResult, NarrowedDiagnostics, check_files_parallel, collect_compiler2_diagnostics, collect_compiler2_diagnostics_narrowed, collect_diagnostics, collect_package_level_diagnostics, }; pub use client_codegen::build_symbol_pool; diff --git a/baml_language/crates/bex_project/src/project.rs b/baml_language/crates/bex_project/src/project.rs index c99aaeeaed0..43f9093f874 100644 --- a/baml_language/crates/bex_project/src/project.rs +++ b/baml_language/crates/bex_project/src/project.rs @@ -224,7 +224,12 @@ pub(crate) fn collect_diagnostic_candidate(guard: &SourceGuard<'_>) -> Diagnosti let mut documents = Vec::new(); let mut has_errors = false; - for file in &source_files { + // Check across worker threads (input order preserved, so per-file zip is + // sound). This runs under the held source guard, so no mutation can race + // the cloned worker handles; the guard is the same exclusion the serial + // loop relied on. + let per_file = baml_project::check_files_parallel(db, &source_files); + for (file, diagnostics) in source_files.iter().zip(per_file) { let file_id = file.file_id(db); let Some(path) = db.file_id_to_path(file_id).cloned() else { continue; @@ -232,7 +237,6 @@ pub(crate) fn collect_diagnostic_candidate(guard: &SourceGuard<'_>) -> Diagnosti let text = file.text(db).clone(); file_texts.insert(file_id, (path.clone(), text)); - let diagnostics = baml_lsp2_actions::check_file(db, *file); has_errors |= diagnostics .iter() .any(|d| d.severity == baml_compiler_diagnostics::Severity::Error);