Vectorize exact MP3 PCM quantization in audio WASM

packages/formats/mp3/rust/Cargo.toml

@@ -12,3 +12,5 @@ path = "src/lib.rs"
 
 [dev-dependencies]
 sha2 = "0.10.9"
+serde_json = "1"
+playsrc-vpk = { path = "../../vpk/rust", default-features = false }

packages/formats/mp3/rust/src/lib.rs

@@ -52,19 +52,73 @@ pub fn decode(bytes: &[u8], max_input: usize, max_samples: usize) -> Result<Deco
         {
             return Err(Error::OutputLimit);
         }
-        for (index, sample) in pcm[..count].iter().enumerate() {
-            output.samples.push(pcm_sample(
-                *sample,
-                (index / usize::from(output.channels)) % 16 == 0,
-            ));
+        let start = output.samples.len();
+        if start + count > output.samples.capacity() {
+            // Reserve a frame without increasing the former per-sample push
+            // capacity bound (power-of-two growth for this i16 buffer).
+            output
+                .samples
+                .reserve_exact((start + count).next_power_of_two() - start);
         }
+        output.samples.resize(start + count, 0);
+        quantize(
+            &pcm[..count],
+            &mut output.samples[start..],
+            usize::from(output.channels),
+        );
     }
     if output.samples.is_empty() {
         return Err(Error::InvalidStream);
     }
     Ok(output)
 }
 
+fn quantize(input: &[f32], output: &mut [i16], channels: usize) {
+    assert_eq!(input.len(), output.len());
+    assert!(channels == 1 || channels == 2);
+    // The first interleaved frame of each synthesis group uses scalar-pair
+    // rounding. The remaining samples are independent nearest-even lanes.
+    for (input, output) in input
+        .chunks(16 * channels)
+        .zip(output.chunks_mut(16 * channels))
+    {
+        let pair = input.len().min(channels);
+        for (sample, value) in input[..pair].iter().zip(&mut output[..pair]) {
+            *value = pcm_sample(*sample, true);
+        }
+        quantize_nearest(&input[pair..], &mut output[pair..]);
+    }
+}
+
+fn quantize_nearest(input: &[f32], output: &mut [i16]) {
+    assert_eq!(input.len(), output.len());
+    #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
+    let at = {
+        use core::arch::wasm32::*;
+        let mut at = 0;
+        while input.len() - at >= 4 {
+            // Four readable f32 inputs and four writable i16 outputs. Neither
+            // load nor store requires vector alignment or touches a tail lane.
+            unsafe {
+                let sample = f32x4_mul(
+                    v128_load(input.as_ptr().add(at).cast()),
+                    f32x4_splat(32768.0),
+                );
+                let rounded = i32x4_trunc_sat_f32x4(f32x4_nearest(sample));
+                let packed = i16x8_narrow_i32x4(rounded, rounded);
+                v128_store64_lane::<0>(packed, output.as_mut_ptr().add(at).cast());
+            }
+            at += 4;
+        }
+        at
+    };
+    #[cfg(not(all(target_arch = "wasm32", target_feature = "simd128")))]
+    let at = 0;
+    for (sample, value) in input[at..].iter().zip(&mut output[at..]) {
+        *value = pcm_sample(*sample, false);
+    }
+}
+
 // minimp3's PC synthesis uses scalar quantization for the pair at positions
 // 0/16 and nearest-even SIMD quantization for the other synthesis samples.
 fn pcm_sample(normalized: f32, scalar_pair: bool) -> i16 {
@@ -102,4 +156,83 @@ mod tests {
         assert_eq!(decode(b"ID3", 2, 100).unwrap_err(), Error::InputLimit);
         assert_eq!(decode(b"ID3", 3, 100).unwrap_err(), Error::InvalidStream);
     }
+
+    #[test]
+    fn synthesis_groups_preserve_every_lane_and_short_tail() {
+        check_synthesis_groups();
+    }
+
+    #[cfg(target_arch = "wasm32")]
+    #[unsafe(no_mangle)]
+    pub extern "C" fn check_wasm_synthesis_groups() {
+        check_synthesis_groups();
+    }
+
+    #[cfg(target_arch = "wasm32")]
+    mod wasm_decode {
+        use std::sync::Mutex;
+        static PCM: Mutex<Vec<i16>> = Mutex::new(Vec::new());
+
+        #[unsafe(no_mangle)]
+        pub extern "C" fn test_input_alloc(length: usize) -> *mut u8 {
+            assert!(length <= 32 * 1024 * 1024);
+            Box::into_raw(vec![0_u8; length].into_boxed_slice()) as *mut u8
+        }
+
+        #[unsafe(no_mangle)]
+        /// # Safety
+        /// Pass one allocation returned by test_input_alloc with its exact length.
+        pub unsafe extern "C" fn test_decode(pointer: *mut u8, length: usize) -> usize {
+            let input =
+                unsafe { Box::from_raw(std::ptr::slice_from_raw_parts_mut(pointer, length)) };
+            let decoded = super::decode(&input, 32 * 1024 * 1024, 32 * 1024 * 1024).unwrap();
+            let mut pcm = PCM.lock().unwrap();
+            *pcm = decoded.samples;
+            pcm.len()
+        }
+
+        #[unsafe(no_mangle)]
+        pub extern "C" fn test_pcm_pointer() -> *const i16 {
+            PCM.lock().unwrap().as_ptr()
+        }
+    }
+
+    fn check_synthesis_groups() {
+        let values = [
+            0.0,
+            -0.0,
+            f32::from_bits(1),
+            -f32::from_bits(1),
+            f32::INFINITY,
+            f32::NEG_INFINITY,
+            f32::NAN,
+            f32::from_bits(0xffc12345),
+            -32767.5 / 32768.0,
+            32766.5 / 32768.0,
+            -1.0 / 32768.0,
+            -2.5 / 32768.0,
+            2.5 / 32768.0,
+            1.0,
+            -1.0,
+        ];
+        for channels in [1, 2] {
+            for offset in 0..values.len() {
+                let input: Vec<_> = (0..97)
+                    .map(|index| values[(index + offset) % values.len()])
+                    .collect();
+                for length in 0..input.len() {
+                    let mut output = vec![12345; length + 2];
+                    quantize(&input[..length], &mut output[1..length + 1], channels);
+                    assert_eq!(output[0], 12345);
+                    assert_eq!(output[length + 1], 12345);
+                    for index in 0..length {
+                        assert_eq!(
+                            output[index + 1],
+                            pcm_sample(input[index], (index / channels) % 16 == 0)
+                        );
+                    }
+                }
+            }
+        }
+    }
 }

packages/formats/mp3/rust/tests/configured.rs

@@ -1,10 +1,56 @@
 use sha2::{Digest, Sha256};
+use std::{
+    fs,
+    io::{Read, Seek, SeekFrom},
+    ops::Range,
+    path::PathBuf,
+};
+
+struct Segments(PathBuf);
+impl playsrc_vpk::SegmentReader for Segments {
+    fn len(&self, index: u32) -> Result<u64, playsrc_vpk::SourceError> {
+        Ok(
+            fs::metadata(self.0.join(format!("tf2_sound_misc_{index:03}.vpk")))
+                .unwrap()
+                .len(),
+        )
+    }
+    fn read(&self, index: u32, range: Range<u64>) -> Result<Vec<u8>, playsrc_vpk::SourceError> {
+        let mut file =
+            fs::File::open(self.0.join(format!("tf2_sound_misc_{index:03}.vpk"))).unwrap();
+        file.seek(SeekFrom::Start(range.start)).unwrap();
+        let mut bytes = vec![0; usize::try_from(range.end - range.start).unwrap()];
+        file.read_exact(&mut bytes).unwrap();
+        Ok(bytes)
+    }
+}
 
 #[test]
-#[ignore = "requires configured cow1.mp3; no game assets are distributed"]
+#[ignore = "requires the exact configured TF2 sound archive; no game assets are distributed"]
 fn configured_mono_matches_public_minimp3_sse_pcm() {
-    let root = std::env::var_os("PLAYSRC_AUDIO_EVIDENCE").expect("explicit evidence directory");
-    let input = std::fs::read(std::path::PathBuf::from(root).join("cow1.mp3")).unwrap();
+    let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../../..");
+    let config: serde_json::Value =
+        serde_json::from_slice(&fs::read(root.join("playsrc.local.json")).unwrap()).unwrap();
+    let contract: serde_json::Value =
+        serde_json::from_slice(&fs::read(root.join("games/tf2/content-build.json")).unwrap())
+            .unwrap();
+    let tf2 = PathBuf::from(config["tf2Dir"].as_str().unwrap());
+    let directory = fs::read(tf2.join("tf2_sound_misc_dir.vpk")).unwrap();
+    assert_eq!(
+        format!("{:x}", Sha256::digest(&directory)),
+        contract["archiveIndexes"]["tf2SoundMisc"].as_str().unwrap()
+    );
+    let archive = playsrc_vpk::parse(
+        &directory,
+        "tf2_sound_misc_dir.vpk",
+        playsrc_vpk::Layout::Split,
+        Default::default(),
+    )
+    .unwrap();
+    let input = archive
+        .read_entry("sound/ambient_mp3/cow1.mp3", &Segments(tf2))
+        .unwrap()
+        .bytes;
     assert_eq!(
         format!("{:x}", Sha256::digest(&input)),
         "6d5029641d1a058b5316d4fd49b7ee923ec6490bb5ce93e40fa25ccaa169aad5"
@@ -14,6 +60,10 @@ fn configured_mono_matches_public_minimp3_sse_pcm() {
         (decoded.sample_rate, decoded.channels, decoded.samples.len()),
         (44100, 1, 73728)
     );
+    assert_eq!(
+        decoded.samples.capacity(),
+        decoded.samples.len().next_power_of_two()
+    );
     let mut digest = Sha256::new();
     for sample in &decoded.samples {
         digest.update(sample.to_le_bytes());
@@ -28,4 +78,8 @@ fn configured_mono_matches_public_minimp3_sse_pcm() {
         playsrc_mp3::decode(&input, input.len(), decoded.samples.len() - 1).unwrap_err(),
         playsrc_mp3::Error::OutputLimit
     );
+    let retained = PathBuf::from(config["sourceCacheDir"].as_str().unwrap())
+        .join("evidence/tf2-wasm-simd-performance/configured");
+    fs::create_dir_all(&retained).unwrap();
+    fs::write(retained.join("cow1.mp3"), &input).unwrap();
 }

packages/presentation/audio/src/playback.ts

@@ -1,4 +1,5 @@
 import { AudioError } from "./error"
+import { compileAudioModule } from "./wasm"
 import type { Listener, NeutralVoice, SoundSource } from "./source"
 
 export type PcmResource = Readonly<{ identity: string; sampleRate: number; numberOfChannels: number; bits: number; length: number; duration: number; loopStartSeconds: number | null }>
@@ -60,8 +61,7 @@ export async function createSourceAudioSystem(context: AudioContext, moduleUrl:
     module = fetch(moduleUrl).then(async response => {
       if (!response.ok) throw new AudioError("BrowserFailure", "Audio module is unavailable")
       const bytes = await response.arrayBuffer()
-      if (bytes.byteLength > 8 * 1024 * 1024) throw new AudioError("Capacity", "Audio module exceeds its bound")
-      return WebAssembly.compile(bytes)
+      return compileAudioModule(bytes)
     })
     moduleCache.set(moduleUrl.href, module)
     void module.catch(() => moduleCache.delete(moduleUrl.href))

packages/presentation/audio/src/wasm.ts

@@ -0,0 +1,12 @@
+import { AudioError } from "./error"
+
+// (module (func (drop (v128.const i32x4 0 0 0 0))))
+// Check standard SIMD without validating the complete audio module twice.
+const SIMD128 = new Uint8Array([0,97,115,109,1,0,0,0,1,4,1,96,0,0,3,2,1,0,10,23,1,21,0,253,12,...Array(16).fill(0),26,11])
+
+/** Reject an unsupported target before compiling or instantiating audio. */
+export function compileAudioModule(bytes: ArrayBuffer): Promise<WebAssembly.Module> {
+  if (bytes.byteLength > 8 * 1024 * 1024) throw new AudioError("Capacity", "Audio module exceeds its bound")
+  if (typeof WebAssembly === "undefined" || !WebAssembly.validate(SIMD128)) throw new AudioError("BrowserFailure", "This browser lacks required standard WebAssembly SIMD128 support")
+  return WebAssembly.compile(bytes)
+}

packages/presentation/audio/tests/wasm.test.ts

@@ -0,0 +1,21 @@
+import { expect, spyOn, test } from "bun:test"
+import { compileAudioModule } from "../src/wasm"
+
+// () -> (), v128.const zero; drop. Standard SIMD, no relaxed instructions.
+const simd = new Uint8Array([0,97,115,109,1,0,0,0,1,4,1,96,0,0,3,2,1,0,10,23,1,21,0,253,12,...Array(16).fill(0),26,11]).buffer
+
+test("audio validates actual SIMD bytecode before compiling, independently of GPU or browser identity", async () => {
+  expect(WebAssembly.validate(simd)).toBe(true)
+  expect(await compileAudioModule(simd)).toBeInstanceOf(WebAssembly.Module)
+  await expect(compileAudioModule(new ArrayBuffer(0))).rejects.toBeInstanceOf(WebAssembly.CompileError)
+  expect(() => compileAudioModule(new ArrayBuffer(8 * 1024 * 1024 + 1))).toThrow("exceeds its bound")
+})
+
+test("an unsupported SIMD target is rejected without compiling the audio module", () => {
+  const validate = spyOn(WebAssembly, "validate").mockReturnValue(false)
+  const compile = spyOn(WebAssembly, "compile")
+  try {
+    expect(() => compileAudioModule(simd)).toThrow("lacks required standard WebAssembly SIMD128")
+    expect(compile).not.toHaveBeenCalled()
+  } finally { validate.mockRestore(); compile.mockRestore() }
+})

playwright.simd-profile.config.ts

@@ -0,0 +1,2 @@
+import { headedProfileConfiguration } from "./tools/playsrc/profile/profile-config"
+export default headedProfileConfiguration({ match: "simd-decoder.profile.ts" })

tools/playsrc/local-jobs.md

@@ -27,6 +27,11 @@ Builds/tests open no browser. The browser stage installs only the Chromium and
 supporting binaries selected by the pinned Playwright package and records the
 executable hash; it never launches a browser or selects a fallback channel.
 
+`simd-decoder` compares the configured MP3 decoder's scalar reference and SIMD
+module in the real headed browser, not gameplay FPS. Prepare its authenticated
+fixtures with `test tools/playsrc/tests/simd-configured.test.ts` at the exact
+checkout commit, then use `PrepareProfile simd-decoder` and ordinary `Run`.
+
 ## Windows consent and ownership
 
 Every Windows `run` above returns a scheduled task identity immediately. The