feat(apple): PyroWave Phase 5 — native Metal decode on Mac / Apple TV / iPad (§4.7)

The Apple client now decodes PyroWave natively on the presenter's own MTLDevice — no MoltenVK, no upstream C++ in the app. Completes and wires up the decoder whose early working-tree snapshot rode along in 9127c346: - MetalWaveletShaders.swift: wavelet_dequant + idwt hand-ported from the vendored GLSL (STORAGE_MODE 0 only; subgroup scans → 32-wide simdgroups; DCShift spec constant → function constant; precision-1 split: fp16 levels 0-1 / fp32 2-4). - MetalWaveletDecoder.swift: Swift reimplementation of push_packet/decode_packet incl. the Phase-4 chunk-aligned window walk (FRAG chains, zeroed missing shards, the >half-blocks partial rule), init_block_meta's block-index space, and the 42-dequant + 13-idwt dispatch structure with encoder-boundary barriers. SOF-dims changes rebuild the size-dependent resources, which is also the mid-stream resize path. Ring of 4 output plane sets on the presenter's queue. - Presenter: pf_frag_planar (3xR8, the planar_csc.frag twin) + renderPlanar with a shared present tail; ReadyFrame carries an image enum (.video | .planar). - Stage2Pipeline: a dedicated PyroWave pump — no VideoToolbox machinery, no keyframe/re-anchor recovery (all-intra; partials render as localized blur by design), newest-frame-index staleness guard for late partials. - Opt-in: "PyroWave (wired LAN)" codec entry (probe-gated, ≈A13 floor via a real kernel-compile probe), selecting it advertises + prefers the codec and forces the session SDR (HDR/10-bit/4:4:4 caps dropped, plan contract). - Core ABI: punktfunk_connection_shard_payload() — the Welcome's negotiated shard payload, needed by native decoders to walk chunk-aligned AUs. - Validation: golden fixtures generated by the host encoder + upstream's own decoder (pyrowave_dump_golden, RTX 5070 Ti); the Metal decode PSNR-matches at 77-88 dB across all planes for dense AND chunk-aligned AUs, and a hole-punched partial still decodes. Parser unit tests cover the window walk, FRAG chains, broken chains, the half-blocks gate, and the block-index layout. Tests: apple 134 green (mac; iOS/tvOS build), host 312 w/ pyrowave on .21, core 148 w/ quic; clippy/fmt clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-15 12:16:43 +02:00
parent a70811043e
commit 739a5f76bf
28 changed files with 826 additions and 81 deletions
@@ -237,10 +237,11 @@ final class AV1Tests: XCTestCase {
        let ready = try XCTUnwrap(frame)
        XCTAssertEqual(ready.ptsNs, 42_000_000)
        XCTAssertFalse(ready.isHDR)
-        XCTAssertEqual(CVPixelBufferGetWidth(ready.pixelBuffer), 320)
-        XCTAssertEqual(CVPixelBufferGetHeight(ready.pixelBuffer), 180)
+        let buffer = try XCTUnwrap(ready.pixelBuffer, "a VT decode delivers a .video frame")
+        XCTAssertEqual(CVPixelBufferGetWidth(buffer), 320)
+        XCTAssertEqual(CVPixelBufferGetHeight(buffer), 180)
        XCTAssertEqual(
-            CVPixelBufferGetPixelFormatType(ready.pixelBuffer),
+            CVPixelBufferGetPixelFormatType(buffer),
            kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange, "SDR AV1 must decode to NV12")
        decoder.reset()
    }
@@ -0,0 +1,292 @@
+// PyroWave Metal decoder tests — two layers:
+//
+// 1. Bitstream/window-walk parser tests (pure CPU): hand-crafted packet streams assert the
+//    exact wire semantics of pyrowave_decoder.cpp's push_packet walk + the Phase-4
+//    chunk-aligned framing (4-byte window prefix, FRAG chains, zeroed missing shards).
+//
+// 2. Golden-frame PSNR tests (Metal GPU): host-encoded fixtures (crates/punktfunk-host
+//    encode/linux/pyrowave.rs `pyrowave_dump_golden`, run on a Vulkan box) decoded by the
+//    Metal port and PSNR-matched against upstream's own decoder output. Float wavelet math is
+//    not bit-exact across implementations (upstream ships precision variants), so the gate is
+//    PSNR, not equality. This is the §4.7 validation oracle for the hand-ported kernels —
+//    the gather/mirror addressing in idwt is the spot most likely to drift.
+
+#if canImport(Metal)
+import Metal
+import XCTest
+
+@testable import PunktfunkKit
+
+final class PyroWaveParserTests: XCTestCase {
+    // 256x144 → aligned 256x160; block space identical to the committed fixtures.
+    private let width = 256
+    private let height = 144
+
+    /// A BitstreamSequenceHeader (START_OF_FRAME) for `width`x`height`, 4:2:0 BT.709 limited.
+    private func sof(totalBlocks: Int, sequence: UInt32 = 1) -> [UInt8] {
+        let word0 =
+            UInt32(width - 1) | (UInt32(height - 1) << 14) | (sequence << 28) | (1 << 31)
+        // code=0 (SOF), chroma=0 (420), primaries/trc/matrix=0 (BT.709), range=1 (LIMITED),
+        // siting=0.
+        let word1 = UInt32(totalBlocks) | (1 << 30)
+        return le32(word0) + le32(word1)
+    }
+
+    /// A minimal coefficient packet: ballot=0 (all 8x8 blocks empty — legal and decodable),
+    /// payload_words=2 (header only).
+    private func packet(blockIndex: Int, sequence: UInt32 = 1) -> [UInt8] {
+        let word0 = UInt32(0) | (2 << 16) | (sequence << 28)
+        let word1 = UInt32(0) | (UInt32(blockIndex) << 8)
+        return le32(word0) + le32(word1)
+    }
+
+    private func le32(_ v: UInt32) -> [UInt8] {
+        [UInt8(v & 0xff), UInt8((v >> 8) & 0xff), UInt8((v >> 16) & 0xff), UInt8(v >> 24)]
+    }
+
+    /// Wrap bodies into `windowSize`-sized windows with the 4-byte used/kind prefix.
+    private func window(_ body: [UInt8], kind: UInt16, size: Int) -> [UInt8] {
+        precondition(body.count + 4 <= size)
+        var out = [UInt8(body.count & 0xff), UInt8(body.count >> 8)]
+        out += [UInt8(kind & 0xff), UInt8(kind >> 8)]
+        out += body
+        out += [UInt8](repeating: 0, count: size - out.count)
+        return out
+    }
+
+    func testLayoutMatchesUpstreamBlockSpace() {
+        // init_block_meta's walk for 256x144 (aligned 256x160): level extents halve from
+        // 128x80; per (comp,level,band) count32 = ceil(ceil(w/8)/4) * ceil(ceil(h/8)/4).
+        let layout = WaveletLayout(width: width, height: height)
+        XCTAssertEqual(layout.alignedWidth, 256)
+        XCTAssertEqual(layout.alignedHeight, 160)
+        XCTAssertEqual(layout.levelWidth(0), 128)
+        XCTAssertEqual(layout.levelHeight(0), 80)
+        XCTAssertEqual(layout.levelWidth(4), 8)
+        XCTAssertEqual(layout.levelHeight(4), 5)
+        // Hand-summed: L4 (8x5 → 1 block) × 3 comps × 4 bands = 12; L3 (16x10 → 1) × 9 = 9;
+        // L2 (32x20 → 1) × 9 = 9; L1 (64x40 → 2x2=4... ) — trust the invariant instead:
+        // every band's count is ceil(w8/4)*ceil(h8/4) and the total is their sum.
+        var expected = 0
+        for level in stride(from: 4, through: 0, by: -1) {
+            let w8 = (layout.levelWidth(level) + 7) / 8
+            let h8 = (layout.levelHeight(level) + 7) / 8
+            let per = ((w8 + 3) / 4) * ((h8 + 3) / 4)
+            for component in 0..<3 {
+                if level == 0 && component != 0 { continue }
+                expected += per * (level == 4 ? 4 : 3)
+            }
+        }
+        XCTAssertEqual(layout.blockCount32, expected)
+        // The finest luma level's stride is its 32-block row width.
+        XCTAssertEqual(layout.blockMeta[0][0][1].stride, (128 + 31) / 32)
+        // Level-0 chroma is not coded in 4:2:0.
+        XCTAssertEqual(layout.blockMeta[1][0][1].offset, -1)
+    }
+
+    func testDenseParseFillsOffsetsAndCountsBlocks() throws {
+        let layout = WaveletLayout(width: width, height: height)
+        var au = sof(totalBlocks: 4)
+        au += packet(blockIndex: 0)
+        au += packet(blockIndex: 3)
+        au += packet(blockIndex: 3) // duplicate — first wins, not double-counted
+        au += packet(blockIndex: layout.blockCount32 - 1)
+        let frame = try XCTUnwrap(
+            WaveletBitstream.parse(au: Data(au), chunkAligned: false, windowSize: 0))
+        XCTAssertEqual(frame.layout.width, width)
+        XCTAssertEqual(frame.totalBlocks, 4)
+        XCTAssertEqual(frame.decodedBlocks, 3)
+        XCTAssertEqual(frame.offsets[0], 0)
+        XCTAssertEqual(frame.offsets[3], 2) // u32 words: each header-only packet is 2 words
+        XCTAssertEqual(frame.offsets[1], UInt32.max)
+        XCTAssertEqual(frame.payload.count, 6)
+        XCTAssertFalse(frame.bt2020)
+        XCTAssertFalse(frame.fullRange) // range bit 1 = LIMITED
+    }
+
+    func testHalfOrFewerBlocksIsDropped() {
+        var au = sof(totalBlocks: 4)
+        au += packet(blockIndex: 0)
+        au += packet(blockIndex: 1)
+        // 2 of 4 decoded = exactly half — upstream requires MORE than half.
+        XCTAssertNil(WaveletBitstream.parse(au: Data(au), chunkAligned: false, windowSize: 0))
+    }
+
+    func testMissingSOFIsDropped() {
+        let au = packet(blockIndex: 0) + packet(blockIndex: 1)
+        XCTAssertNil(WaveletBitstream.parse(au: Data(au), chunkAligned: false, windowSize: 0))
+    }
+
+    func testTruncatedPacketIsRejected() {
+        var au = sof(totalBlocks: 1)
+        // Claims 4 payload words but only the 8-byte header follows.
+        let word0 = UInt32(0) | (4 << 16) | (1 << 28)
+        au += le32(word0) + le32(0)
+        XCTAssertNil(WaveletBitstream.parse(au: Data(au), chunkAligned: false, windowSize: 0))
+    }
+
+    func testWindowWalkPackedFragAndMissingShard() throws {
+        let size = 64
+        // Window 1: SOF + one packet, PACKED. Window 2: a FRAG chain carrying one packet split
+        // across two windows. Window 3: all zeros (a lost shard of a partial frame). Window 4:
+        // a PACKED packet — the chain break must not eat it.
+        let fragPacket = packet(blockIndex: 2)
+        var au = window(sof(totalBlocks: 3) + packet(blockIndex: 0), kind: 0, size: size)
+        au += window(Array(fragPacket[0..<5]), kind: 1, size: size)
+        au += window(Array(fragPacket[5...]), kind: 3, size: size)
+        au += [UInt8](repeating: 0, count: size) // missing shard
+        au += window(packet(blockIndex: 1), kind: 0, size: size)
+        let frame = try XCTUnwrap(
+            WaveletBitstream.parse(au: Data(au), chunkAligned: true, windowSize: size))
+        XCTAssertEqual(frame.decodedBlocks, 3)
+        XCTAssertEqual(frame.offsets[0], 0)
+        XCTAssertEqual(frame.offsets[2], 2)
+        XCTAssertEqual(frame.offsets[1], 4)
+    }
+
+    func testBrokenFragChainIsDiscarded() throws {
+        let size = 64
+        let fragPacket = packet(blockIndex: 2)
+        var au = window(sof(totalBlocks: 1) + packet(blockIndex: 0), kind: 0, size: size)
+        au += window(Array(fragPacket[0..<5]), kind: 1, size: size)
+        au += [UInt8](repeating: 0, count: size) // the chain's middle shard was lost
+        au += window(Array(fragPacket[5...]), kind: 3, size: size) // dangling LAST — dropped
+        let frame = try XCTUnwrap(
+            WaveletBitstream.parse(au: Data(au), chunkAligned: true, windowSize: size))
+        XCTAssertEqual(frame.decodedBlocks, 1)
+        XCTAssertEqual(frame.offsets[2], UInt32.max)
+    }
+}
+
+/// Golden-frame decode against the committed host-encoder fixtures. Skipped when the machine
+/// has no Metal device (headless CI) — everywhere else this is the hand-ported kernels' guard.
+final class PyroWaveGoldenTests: XCTestCase {
+    private static let fixtureDir = "PyroWaveFixtures"
+
+    private func fixture(_ name: String) throws -> Data {
+        let url = try XCTUnwrap(
+            Bundle.module.url(
+                forResource: name, withExtension: "bin", subdirectory: Self.fixtureDir),
+            "missing fixture \(name).bin — regenerate with pyrowave_dump_golden")
+        return try Data(contentsOf: url)
+    }
+
+    /// Completion box — the decode callback lands on a Metal thread.
+    private final class ResultBox: @unchecked Sendable {
+        let lock = NSLock()
+        var planes: WaveletPlanes?
+    }
+
+    /// Decode `au` synchronously and read all three planes back to CPU bytes.
+    private func decode(
+        au: Data, chunkAligned: Bool, windowSize: Int
+    ) throws -> (y: [UInt8], cb: [UInt8], cr: [UInt8]) {
+        let device = try XCTUnwrap(MTLCreateSystemDefaultDevice())
+        let queue = try XCTUnwrap(device.makeCommandQueue())
+        let decoder = try XCTUnwrap(MetalWaveletDecoder(device: device, queue: queue))
+        let done = expectation(description: "decode completes")
+        let box = ResultBox()
+        let submitted = decoder.decode(
+            au: au, chunkAligned: chunkAligned, windowSize: windowSize
+        ) { planes in
+            box.lock.lock()
+            box.planes = planes
+            box.lock.unlock()
+            done.fulfill()
+        }
+        XCTAssertTrue(submitted, "the fixture AU must parse")
+        wait(for: [done], timeout: 10)
+        box.lock.lock()
+        let result = box.planes
+        box.lock.unlock()
+        let planes = try XCTUnwrap(result, "the GPU pass must complete without error")
+        return (
+            try readback(planes.y, device: device, queue: queue),
+            try readback(planes.cb, device: device, queue: queue),
+            try readback(planes.cr, device: device, queue: queue)
+        )
+    }
+
+    private func readback(
+        _ texture: MTLTexture, device: MTLDevice, queue: MTLCommandQueue
+    ) throws -> [UInt8] {
+        let bytesPerRow = texture.width
+        let length = bytesPerRow * texture.height
+        let buffer = try XCTUnwrap(device.makeBuffer(length: length, options: .storageModeShared))
+        let cmd = try XCTUnwrap(queue.makeCommandBuffer())
+        let blit = try XCTUnwrap(cmd.makeBlitCommandEncoder())
+        blit.copy(
+            from: texture, sourceSlice: 0, sourceLevel: 0,
+            sourceOrigin: MTLOrigin(x: 0, y: 0, z: 0),
+            sourceSize: MTLSize(width: texture.width, height: texture.height, depth: 1),
+            to: buffer, destinationOffset: 0, destinationBytesPerRow: bytesPerRow,
+            destinationBytesPerImage: length)
+        blit.endEncoding()
+        cmd.commit()
+        cmd.waitUntilCompleted()
+        return [UInt8](UnsafeRawBufferPointer(start: buffer.contents(), count: length))
+    }
+
+    private func psnr(_ a: [UInt8], _ b: [UInt8]) -> Double {
+        precondition(a.count == b.count)
+        var sse = 0.0
+        for i in 0..<a.count {
+            let d = Double(a[i]) - Double(b[i])
+            sse += d * d
+        }
+        if sse == 0 { return .infinity }
+        let mse = sse / Double(a.count)
+        return 10 * log10(255.0 * 255.0 / mse)
+    }
+
+    private func assertMatchesReference(
+        _ decoded: (y: [UInt8], cb: [UInt8], cr: [UInt8]), prefix: String,
+        file: StaticString = #filePath, line: UInt = #line
+    ) throws {
+        for (name, plane, ref) in [
+            ("y", decoded.y, try fixture("\(prefix)-y")),
+            ("cb", decoded.cb, try fixture("\(prefix)-cb")),
+            ("cr", decoded.cr, try fixture("\(prefix)-cr")),
+        ] {
+            XCTAssertEqual(plane.count, ref.count, file: file, line: line)
+            let db = psnr(plane, [UInt8](ref))
+            print("pyrowave golden \(prefix) \(name): \(db) dB")
+            // The Metal port and upstream's decoder run the same math at the same precision
+            // tier; residual differences are float rounding + the gather/mirror edge handling.
+            // Well-matched ports measure ≫50 dB; 45 catches a real divergence long before it
+            // is visible.
+            XCTAssertGreaterThan(db, 45.0, "plane PSNR \(db) dB", file: file, line: line)
+        }
+    }
+
+    func testDenseGoldenFrame() throws {
+        try XCTSkipIf(!MetalWaveletDecoder.supported, "no capable Metal device")
+        let au = try fixture("au-dense")
+        let decoded = try decode(au: au, chunkAligned: false, windowSize: 0)
+        try assertMatchesReference(decoded, prefix: "ref-dense")
+    }
+
+    func testChunkAlignedGoldenFrame() throws {
+        try XCTSkipIf(!MetalWaveletDecoder.supported, "no capable Metal device")
+        let au = try fixture("au-chunked")
+        let decoded = try decode(au: au, chunkAligned: true, windowSize: 1408)
+        try assertMatchesReference(decoded, prefix: "ref-chunked")
+    }
+
+    /// Phase-4 partial delivery: zero a mid-AU window (a lost shard) — the frame must still
+    /// decode (blocks > half) and stay recognizably the same picture (holes reconstruct as
+    /// localized blur, not garbage).
+    func testPartialFrameStillDecodes() throws {
+        try XCTSkipIf(!MetalWaveletDecoder.supported, "no capable Metal device")
+        var au = try fixture("au-chunked")
+        let windows = au.count / 1408
+        try XCTSkipIf(windows < 3, "fixture too small to punch a hole in")
+        let hole = (windows / 2) * 1408
+        au.replaceSubrange(hole..<(hole + 1408), with: [UInt8](repeating: 0, count: 1408))
+        let decoded = try decode(au: au, chunkAligned: true, windowSize: 1408)
+        let ref = try fixture("ref-chunked-y")
+        let db = psnr(decoded.y, [UInt8](ref))
+        XCTAssertGreaterThan(db, 25.0, "lossy frame should still resemble the source (\(db) dB)")
+    }
+}
+#endif
@@ -47,18 +47,21 @@ final class Stage444Tests: XCTestCase {
        box.lock.lock(); let frame = box.frame; let error = box.error; box.lock.unlock()
        XCTAssertNil(error.map { "decode error \($0)" })
        let ready = try XCTUnwrap(frame, "a 4:4:4 ReadyFrame must be delivered")
-        XCTAssertEqual(CVPixelBufferGetWidth(ready.pixelBuffer), 256)
-        XCTAssertEqual(CVPixelBufferGetHeight(ready.pixelBuffer), 256)
-        let pf = CVPixelBufferGetPixelFormatType(ready.pixelBuffer)
+        guard case .video(let buffer, let isHDR) = ready.image else {
+            return XCTFail("a VideoToolbox decode must deliver a .video frame")
+        }
+        XCTAssertEqual(CVPixelBufferGetWidth(buffer), 256)
+        XCTAssertEqual(CVPixelBufferGetHeight(buffer), 256)
+        let pf = CVPixelBufferGetPixelFormatType(buffer)
        XCTAssertTrue(
            pf == kCVPixelFormatType_444YpCbCr8BiPlanarVideoRange
                || pf == kCVPixelFormatType_444YpCbCr8BiPlanarFullRange,
            "expected a biplanar 4:4:4 8-bit buffer, got \(fourCCString(pf))")
-        XCTAssertFalse(ready.isHDR, "an 8-bit BT.709 4:4:4 stream is SDR")
+        XCTAssertFalse(isHDR, "an 8-bit BT.709 4:4:4 stream is SDR")
        // The chroma plane (plane 1) must be FULL resolution for 4:4:4 (vs half for 4:2:0) — this is
        // what lets the unchanged shader sample chroma at the luma UV.
-        XCTAssertEqual(CVPixelBufferGetWidthOfPlane(ready.pixelBuffer, 1), 256)
-        XCTAssertEqual(CVPixelBufferGetHeightOfPlane(ready.pixelBuffer, 1), 256)
+        XCTAssertEqual(CVPixelBufferGetWidthOfPlane(buffer, 1), 256)
+        XCTAssertEqual(CVPixelBufferGetHeightOfPlane(buffer, 1), 256)
    }

    private func fourCCString(_ t: OSType) -> String {
@@ -99,8 +99,9 @@ final class VideoToolboxRoundTripTests: XCTestCase {
        box.lock.unlock()
        XCTAssertNil(error.map { "decode error \($0)" })
        let ready = try XCTUnwrap(frame, "the async output callback must deliver a ReadyFrame")
-        XCTAssertEqual(CVPixelBufferGetWidth(ready.pixelBuffer), width)
-        XCTAssertEqual(CVPixelBufferGetHeight(ready.pixelBuffer), height)
+        let buffer = try XCTUnwrap(ready.pixelBuffer, "a VT decode delivers a .video frame")
+        XCTAssertEqual(CVPixelBufferGetWidth(buffer), width)
+        XCTAssertEqual(CVPixelBufferGetHeight(buffer), height)
        XCTAssertEqual(ready.ptsNs, 42_000_000, "pts round-trips through the decoder")
        XCTAssertEqual(
            ready.receivedNs, 41_000_000, "receivedNs round-trips through the frame refcon")