cli: glibc's allocator carries a lot of overhead on the graphs the be…

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"

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

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

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.<MALLCTL_ARENAS_ALL>.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 {

cli/src/params.rs

@@ -840,6 +840,7 @@ impl Parameters {
                 Ok(m.nested_models(node)[0].1.downcast_ref::<TypedModel>().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,

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

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.