1use 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}