diff --git a/Cargo.lock b/Cargo.lock index 6d43c59df8..33aa06f711 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4861,6 +4861,26 @@ dependencies = [ "zune-jpeg", ] +[[package]] +name = "tikv-jemalloc-sys" +version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd8aa5b2ab86a2cefa406d889139c162cbb230092f7d1d7cbc1716405d852a3b" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0359b4327f954e0567e69fb191cf1436617748813819c94b8cd4a431422d053a" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.55" @@ -5191,6 +5211,8 @@ dependencies = [ "serde", "serde_json", "tar", + "tikv-jemalloc-sys", + "tikv-jemallocator", "time", "toml", "tract-core", diff --git a/Cargo.toml b/Cargo.toml index 8002c47153..07c7def59b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -206,6 +206,8 @@ rand = "0.10" rand_distr = "0.6" rayon = "1.10" readings-probe = "0.1.8" +tikv-jemalloc-sys = "0.6" +tikv-jemallocator = "0.6" regex = "1.5.4" reqwest = { version = "0.13", features = [ "blocking", "rustls-no-provider" ], default-features = false } ron = "0.12" diff --git a/ci/cross.sh b/ci/cross.sh index 882938a715..49ff00a86c 100755 --- a/ci/cross.sh +++ b/ci/cross.sh @@ -230,7 +230,11 @@ case "$PLATFORM" in # riscv64 musl defaults to dynamic linking (unlike the aarch64/armv7 musl # targets here); force a static binary so it runs on the glibc boards, which # have no musl loader. - export CARGO_TARGET_RISCV64GC_UNKNOWN_LINUX_MUSL_RUSTFLAGS="-C target-feature=+crt-static" + # jemalloc's ffs lowers to a __ffsdi2 libgcc call on rv64gc, which has no + # count-trailing-zeros outside Zbb, and rust's compiler_builtins carries the + # rest of the *di2 family but not that one: name libgcc for the -nodefaultlibs + # link. + export CARGO_TARGET_RISCV64GC_UNKNOWN_LINUX_MUSL_RUSTFLAGS="-C target-feature=+crt-static -C link-arg=-lgcc" ;; # RVV is vector-length agnostic and the mmm kernels are gated on the # hart's VLEN, so the two entries below differ only in vlen: 256 is diff --git a/cli/Cargo.toml b/cli/Cargo.toml index f2b993569a..6ea0e33ffa 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -96,6 +96,10 @@ tract-cuda = { workspace = true, optional = true } [target.'cfg(target_os = "linux")'.dependencies] rustix = { version = "1", features = ["fs"] } +[target.'cfg(not(any(target_family = "wasm", target_os = "windows", target_vendor = "apple")))'.dependencies] +tikv-jemalloc-sys = { workspace = true, optional = true } +tikv-jemallocator = { workspace = true, optional = true } + [features] default = [ "onnx", @@ -121,7 +125,8 @@ tf = ["dep:tract-tensorflow", "tract-libcli/hir"] tflite = ["dep:tract-tflite"] transformers = ["dep:tract-transformers", "tract-libcli/transformers", "tract-onnx?/transformers"] multithread-mm = ["tract-linalg/multithread-mm"] -bench-suite = ["dep:toml", "dep:tar", "dep:minijinja", "dep:time", "multithread-mm"] +bench-suite = ["dep:toml", "dep:tar", "dep:minijinja", "dep:time", "multithread-mm", "jemalloc"] +jemalloc = ["dep:tikv-jemalloc-sys", "dep:tikv-jemallocator"] # Marker feature implicitly enabled by every cuda-XXXXX selector below. # Use to gate cudarc-specific code paths in tract-cli. Not meant to be diff --git a/cli/src/main.rs b/cli/src/main.rs index e9711d0c82..5bc9c727e1 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -59,7 +59,48 @@ use tract_linalg::WeightType; use tract_linalg::block_quant::Q4_0; use tract_linalg::mmm::MatMatMul; -readings_probe::instrumented_allocator!(); +// Darwin's allocator does not have the glibc mmap/munmap-per-eval pathology this works +// around, and jemalloc measured slower there, so apple targets keep the system allocator. +#[cfg(all( + feature = "jemalloc", + not(any(target_family = "wasm", target_os = "windows", target_vendor = "apple")) +))] +mod allocator { + readings_probe::wrap_global_allocator!(tikv_jemallocator::Jemalloc); + + // jemalloc keeps purged extents mapped and dirty under its default `retain`, and some + // kernels never take those pages back: the resident set then reports the process peak + // rather than what the model holds. + #[unsafe(export_name = "_rjem_malloc_conf")] + #[allow(non_upper_case_globals)] + pub static malloc_conf: &[u8] = b"retain:false\0"; + + /// Return the pages jemalloc holds dirty for reuse to the OS, so a resident-set + /// reading taken right after reports what the process holds live rather than what + /// it freed within the last decay period. + pub fn purge() { + // SAFETY: `arena..purge` takes and returns no value. + unsafe { + tikv_jemalloc_sys::mallctl( + c"arena.4096.purge".as_ptr(), + std::ptr::null_mut(), + std::ptr::null_mut(), + std::ptr::null_mut(), + 0, + ); + } + } +} + +#[cfg(not(all( + feature = "jemalloc", + not(any(target_family = "wasm", target_os = "windows", target_vendor = "apple")) +)))] +mod allocator { + readings_probe::instrumented_allocator!(); + + pub fn purge() {} +} fn info_usage(stage: &str, probe: Option<&Probe>) { if let Some(mon) = probe { diff --git a/cli/src/params.rs b/cli/src/params.rs index d4b4689101..dbb53ce059 100644 --- a/cli/src/params.rs +++ b/cli/src/params.rs @@ -840,6 +840,7 @@ impl Parameters { Ok(m.nested_models(node)[0].1.downcast_ref::().unwrap().clone()) }); } + crate::allocator::purge(); stage!("before-optimize", typed_model -> typed_model, Ok); Ok((typed_model.clone().unwrap(), reference_model)) } @@ -1092,6 +1093,7 @@ impl Parameters { }; info!("Model ready"); + crate::allocator::purge(); info_usage("model ready", probe); Ok(Parameters { graph, diff --git a/doc/README.md b/doc/README.md index 2aae79fd72..b12e505d30 100644 --- a/doc/README.md +++ b/doc/README.md @@ -11,6 +11,7 @@ that is harder to derive from reading the source. * [`graph.md`](graph.md) — Graph, Node, Outlet, Fact, model pipeline * [`op.md`](op.md) — anatomy of an Op (`Op` / `EvalOp` / `TypedOp` / `InferenceOp`) * [`cli-recipe.md`](cli-recipe.md) — `tract` command-line cookbook +* [`allocator.md`](allocator.md) — which allocator to link, and how to configure it * [`pgo.md`](pgo.md) — profile-guided optimization build recipe * [`kernel-notes.md`](kernel-notes.md) — tract-linalg kernels and debugging * [`pgo-bolt.md`](pgo-bolt.md) — profile-guided optimization and BOLT release builds diff --git a/doc/allocator.md b/doc/allocator.md new file mode 100644 index 0000000000..34c7367101 --- /dev/null +++ b/doc/allocator.md @@ -0,0 +1,100 @@ +# Allocator choice + +tract allocates on every eval. A plan hands each op its output tensor fresh, +so a streaming graph mallocs and frees a handful of intermediates per pulse, +and a load does one large burst while parsing and optimizing. The allocator is +therefore part of the engine's cost, and it belongs to the application: tract +is a library, so nothing below is something the engine can pick for you. + +## Use jemalloc + +glibc's allocator is the expensive one on the graphs tract runs. On small +pulsed graphs — the ones whose per-turn cost is mostly bookkeeping rather than +a steady-state matmul — jemalloc is worth 5 to 9% of eval time on aarch64 and +riscv64, and about half that on x86_64. It also cuts the minor-fault count of a +load by an order of magnitude: glibc can settle into a regime where a per-eval +intermediate is mapped and unmapped on every pass. musl's allocator has the +same failure mode at a smaller grain: a process can land in a regime where the +same small slot groups are unmapped and remapped on every eval. + +Darwin's allocator has neither problem, and jemalloc measured slower there: +keep the system allocator on Apple targets. + +The gain is not reachable by tuning glibc. Raising `M_MMAP_THRESHOLD` and +`M_TRIM_THRESHOLD` (or their `MALLOC_MMAP_THRESHOLD_` / `MALLOC_TRIM_THRESHOLD_` +environment forms) removes most of the fault churn and none of the time: what +costs is the allocator's own fast path, not the mapping. + +```toml +[dependencies] +tikv-jemallocator = "0.6" +``` + +```rust +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; +``` + +## Set `retain:false` + +jemalloc's `opt.retain` defaults to true on 64-bit targets: a purged extent +keeps both its mapping and its dirty pages, so the allocator can reuse the +address range without asking the kernel again. Whether those pages ever come +back is up to the kernel, and on some of them they do not. A graph holding +68 MB of live tensors sat in a 338 MB resident set on a 5.4 aarch64 kernel, +while the same binary on a 4.9 kernel of the same SoC stayed at 115 MB. The +pages are dirty and anonymous, not lazily freed, so nothing reclaims them +under memory pressure short of swap. + +Turning retain off restores a resident set that tracks the live one, and costs +nothing in eval time — the extra mapping work lands in the load path, where it +is worth about 8% of the time to a ready model on a big graph. It is already +jemalloc's own default on 32-bit targets. + +Bake it into the binary, so a deployment cannot lose it by clearing the +environment: + +```rust +#[allow(non_upper_case_globals)] +#[unsafe(export_name = "_rjem_malloc_conf")] +pub static malloc_conf: &[u8] = b"retain:false\0"; +``` + +The symbol name carries tikv-jemalloc-sys' `_rjem_` prefix; the unprefixed +`malloc_conf` is what a build without that prefix wants. The same string can +be passed at run time as `_RJEM_MALLOC_CONF` for a quick check. + +Do not reach for `dirty_decay_ms:0` / `muzzy_decay_ms:0` on top of it. Every +free then purges immediately, so a pulsed graph pays the mapping churn again on +each eval, and most of what jemalloc brings is gone. + +## Resident set right after load + +jemalloc keeps freed pages mapped for its dirty decay period (10 s by default) +so it can hand them out again without a fault. A load frees most of what it +allocates within that period, so the resident set measured right after a model +is ready can be a large multiple of what it holds on a small model. It settles +once the decay runs. An application that wants the settled figure at once can +purge every arena after loading: + +```rust +unsafe { + tikv_jemalloc_sys::mallctl( + c"arena.4096.purge".as_ptr(), + std::ptr::null_mut(), + std::ptr::null_mut(), + std::ptr::null_mut(), + 0, + ); +} +``` + +`4096` is `MALLCTL_ARENAS_ALL`. + +## Benching + +`tract`'s own CLI links jemalloc with this configuration under its `jemalloc` +feature, which `bench-suite` turns on (Apple targets excepted), so the bench +numbers describe the configuration above rather than a glibc or musl build. The +CLI purges right before its load and model-ready readings, so the RSS it records +there is the settled figure.