Skip to main content

rlx_tsac/
perf.rs

1//! Host RLX perf bench vs Docker reference binaries (Bellard + tsac-ng).
2
3use crate::audio::{self, PcmCompare, prepare_tsac_wav};
4use crate::codec::{TsacBackendKind, TsacCodec, TsacOptions};
5use crate::docker_ref::{DockerRefOptions, RefRoundtrip, run_docker_ref_roundtrip};
6use crate::download::weights_available;
7use crate::parity::EngineRoundtrip;
8use anyhow::{Result, bail};
9use rlx_runtime::Device;
10use std::path::{Path, PathBuf};
11use std::time::Instant;
12
13#[derive(Debug, Clone)]
14pub struct PerfOptions {
15    pub quality: u8,
16    pub fast: bool,
17    pub native_device: Device,
18    pub docker: DockerRefOptions,
19    pub min_correlation: f32,
20    pub max_mse: f32,
21}
22
23impl Default for PerfOptions {
24    fn default() -> Self {
25        Self {
26            quality: 9,
27            fast: true,
28            native_device: Device::Cpu,
29            docker: DockerRefOptions::default(),
30            min_correlation: env_f32("RLX_TSAC_PARITY_MIN_CORR", 0.92),
31            max_mse: env_f32("RLX_TSAC_PARITY_MAX_MSE", 0.0025),
32        }
33    }
34}
35
36#[derive(Debug, Clone)]
37pub struct PerfBenchReport {
38    pub input_wav: PathBuf,
39    pub prepared_wav: PathBuf,
40    pub channels: u16,
41    pub rlx: EngineRoundtrip,
42    pub bellard: RefRoundtrip,
43    pub tsac_ng: RefRoundtrip,
44    pub rlx_vs_bellard: PcmCompare,
45    pub rlx_vs_tsac_ng: PcmCompare,
46}
47
48impl PerfBenchReport {
49    pub fn print_summary(&self, opts: &PerfOptions) {
50        eprintln!(
51            "\n=== TSAC perf (RLX host vs Docker refs) ===\n\
52             input: {}\n\
53             prepared @ 44.1 kHz: {} ({} ch)\n\
54             quality={} fast={} device={:?}\n",
55            self.input_wav.display(),
56            self.prepared_wav.display(),
57            self.channels,
58            opts.quality,
59            opts.fast,
60            opts.native_device,
61        );
62        print_engine("rlx (host)", &self.rlx);
63        print_ref("bellard (docker)", &self.bellard);
64        print_ref("tsac-ng (docker)", &self.tsac_ng);
65        print_compare("RLX vs bellard PCM", &self.rlx_vs_bellard, opts);
66        print_compare("RLX vs tsac-ng PCM", &self.rlx_vs_tsac_ng, opts);
67    }
68}
69
70pub fn bench_perf(
71    install_dir: impl AsRef<Path>,
72    in_wav: impl AsRef<Path>,
73    opts: &PerfOptions,
74) -> Result<PerfBenchReport> {
75    let install_dir = install_dir.as_ref();
76    if !weights_available(install_dir) {
77        bail!(
78            "TSAC weights missing under {} — run `just fetch-tsac`",
79            install_dir.display()
80        );
81    }
82    let in_wav = in_wav.as_ref().to_path_buf();
83    let tag = std::process::id();
84    let work = std::env::temp_dir().join(format!("rlx-tsac-perf-{tag}"));
85    std::fs::create_dir_all(&work)?;
86
87    let prepared = work.join("input_44100.wav");
88    let channels = prepare_tsac_wav(&in_wav, &prepared)?;
89
90    let bellard = run_docker_ref_roundtrip(
91        &opts.docker,
92        "bellard",
93        &prepared,
94        &work,
95        opts.quality,
96        opts.fast,
97    )?;
98    let tsac_ng = run_docker_ref_roundtrip(
99        &opts.docker,
100        "tsac-ng",
101        &prepared,
102        &work,
103        opts.quality,
104        opts.fast,
105    )?;
106
107    let native_opts = TsacOptions {
108        quality: Some(opts.quality),
109        fast: opts.fast,
110        verbose: false,
111        separate_stereo: false,
112        channels: None,
113        device: opts.native_device,
114        backend: TsacBackendKind::Native,
115    };
116    let codec = TsacCodec::open_with_options(install_dir, native_opts)?;
117    let rlx = run_rlx_roundtrip(&codec, "rlx", &prepared, &work)?;
118
119    let bellard_pcm = audio::load_pcm_from_wav(&bellard.wav_path)?;
120    let tsac_ng_pcm = audio::load_pcm_from_wav(&tsac_ng.wav_path)?;
121    let rlx_vs_bellard = PcmCompare::compare(&rlx.pcm, &bellard_pcm);
122    let rlx_vs_tsac_ng = PcmCompare::compare(&rlx.pcm, &tsac_ng_pcm);
123
124    Ok(PerfBenchReport {
125        input_wav: in_wav,
126        prepared_wav: prepared,
127        channels,
128        rlx,
129        bellard,
130        tsac_ng,
131        rlx_vs_bellard,
132        rlx_vs_tsac_ng,
133    })
134}
135
136fn run_rlx_roundtrip(
137    codec: &TsacCodec,
138    label: &str,
139    prepared_wav: &Path,
140    work: &Path,
141) -> Result<EngineRoundtrip> {
142    let tsac_path = work.join(format!("{label}.tsac"));
143    let wav_path = work.join(format!("{label}_roundtrip.wav"));
144    let encode = codec.encode(prepared_wav, &tsac_path)?;
145    let t0 = Instant::now();
146    codec.decode(&tsac_path, &wav_path)?;
147    let decode_ms = t0.elapsed().as_secs_f64() * 1000.0;
148    let pcm = audio::load_pcm_from_wav(&wav_path)?;
149    Ok(EngineRoundtrip {
150        encode,
151        decode_ms,
152        pcm,
153        tsac_path,
154        wav_path,
155    })
156}
157
158fn print_engine(label: &str, rt: &EngineRoundtrip) {
159    eprintln!(
160        "{label}: encode {:.1} ms, decode {:.1} ms, {} bytes compressed, {} samples",
161        rt.encode.encode_ms,
162        rt.decode_ms,
163        rt.encode.output_bytes,
164        rt.pcm.len()
165    );
166}
167
168fn print_ref(label: &str, rt: &RefRoundtrip) {
169    eprintln!(
170        "{label}: encode {:.1} ms, decode {:.1} ms, {} bytes compressed",
171        rt.encode_ms, rt.decode_ms, rt.output_bytes
172    );
173}
174
175fn print_compare(label: &str, cmp: &PcmCompare, opts: &PerfOptions) {
176    let ok = cmp.passes(opts.min_correlation, opts.max_mse);
177    eprintln!(
178        "{label}: corr={:.4} mse={:.6} max_abs={:.5} n={} [{}]",
179        cmp.correlation,
180        cmp.mse,
181        cmp.max_abs,
182        cmp.samples,
183        if ok { "ok" } else { "over tolerance" }
184    );
185}
186
187fn env_f32(name: &str, default: f32) -> f32 {
188    std::env::var(name)
189        .ok()
190        .and_then(|s| s.parse().ok())
191        .unwrap_or(default)
192}