diff --git a/clients/apple/Sources/PunktfunkKit/Video/MetalWaveletDecoder.swift b/clients/apple/Sources/PunktfunkKit/Video/MetalWaveletDecoder.swift new file mode 100644 index 00000000..a440cb19 --- /dev/null +++ b/clients/apple/Sources/PunktfunkKit/Video/MetalWaveletDecoder.swift @@ -0,0 +1,604 @@ +// PyroWave native Metal decoder — the Apple twin of pf-client-core's Vulkan decoder +// (crates/pf-client-core/src/video_pyrowave.rs), reimplemented on the presenter's own MTLDevice +// so decode + CSC + present share one device with zero interop (design/pyrowave-codec-plan.md +// §4.7). No upstream C/C++ ships in the app: the bitstream parse below reimplements +// pyrowave_decoder.cpp's push_packet/decode_packet walk, and the two compute kernels +// (MetalWaveletShaders.swift) are hand-ported from the vendored GLSL. The §4.2 upstream pin +// covers this hand-port: a vendored bump means re-diffing two decode shaders and the two 8-byte +// header structs, and it is already a protocol-version event. +// +// Wire shape (all fixed by the host encoder, punktfunk-host encode/linux/pyrowave.rs): +// • One AU = one frame = a self-delimiting stream of packets. Each packet is one 32x32 +// coefficient block for one (component, level, band), self-sized by its 8-byte +// BitstreamHeader; a per-frame START_OF_FRAME sequence header carries dims + total block +// count + the VUI bits (chroma 4:2:0, BT.709/BT.2020, limited/full). +// • With `USER_FLAG_CHUNK_ALIGNED` (Phase 4) the AU is a whole number of `shard_payload`-sized +// windows, each 4-byte-prefixed (used-len u16 LE + kind u16 LE): kind 0 = whole packets, +// 1/2/3 = FRAG chain for a packet bigger than one window. A missing shard of a partial frame +// arrives as an all-zero window (used = 0) → skipped, its blocks reconstruct as zeros +// (localized blur, the Phase-4 design intent). The reassembler enables partial delivery +// core-side automatically for PyroWave sessions. +// • Decode acceptance mirrors upstream decode_is_ready(allow_partial=true): a frame with no +// SOF or with no more than half its blocks is dropped rather than decoded to garbage. +// +// GPU structure per frame (mirroring pyrowave_decoder.cpp's barriers): one concurrent compute +// encoder with all ~42 dequant dispatches (each writes a distinct band layer — no intra-stage +// hazards), then one concurrent encoder per iDWT level (5) — encoder boundaries provide the +// write→sampled-read synchronization the Vulkan version expresses as pipeline barriers. The +// output is a ring of 4 plane sets (Y full-res + Cb/Cr half-res R8Unorm); ring depth plus +// same-queue hazard tracking keeps a set alive while the presenter still samples it (the same +// scheme as the Vulkan client's ring). + +#if canImport(Metal) +import Foundation +import Metal +import os + +private let waveletLog = Logger(subsystem: "io.unom.punktfunk", category: "pyrowave") + +/// The per-(component, level, band) 32x32-block table — the exact Swift port of +/// `WaveletBuffers::init_block_meta` (pyrowave_common.cpp): the walk order (level 4→0, +/// component 0→2 skipping level-0 chroma in 4:2:0, band (level==4 ? 0 : 1)→3) DEFINES the +/// global `block_index` space the wire packets address, so it must match the encoder exactly. +struct WaveletLayout { + static let decompositionLevels = 5 + static let alignment = 32 + static let minimumImageSize = 128 + + let width: Int + let height: Int + let alignedWidth: Int + let alignedHeight: Int + /// blockMeta[component][level][band] = (blockOffset32x32, blockStride32x32); -1 offset = + /// band not coded (level-0 chroma in 4:2:0). + let blockMeta: [[[(offset: Int, stride: Int)]]] + let blockCount32: Int + + /// Band-image extent at `level` — mip `level` of the (aligned/2)-sized coefficient image. + /// Exact halving: the aligned dims are 32-aligned, so /2 is 16-aligned and survives 4 shifts. + func levelWidth(_ level: Int) -> Int { (alignedWidth / 2) >> level } + func levelHeight(_ level: Int) -> Int { (alignedHeight / 2) >> level } + + init(width: Int, height: Int) { + self.width = width + self.height = height + let align = { (v: Int) in + max((v + Self.alignment - 1) & ~(Self.alignment - 1), Self.minimumImageSize) + } + alignedWidth = align(width) + alignedHeight = align(height) + + var meta = [[[(offset: Int, stride: Int)]]]( + repeating: [[(offset: Int, stride: Int)]]( + repeating: [(offset: Int, stride: Int)](repeating: (-1, 0), count: 4), + count: Self.decompositionLevels), + count: 3) + var count32 = 0 + let aw = alignedWidth + let ah = alignedHeight + for level in stride(from: Self.decompositionLevels - 1, through: 0, by: -1) { + for component in 0..<3 { + if level == 0 && component != 0 { continue } // 4:2:0: no top-level chroma + for band in (level == Self.decompositionLevels - 1 ? 0 : 1)..<4 { + let levelW = (aw / 2) >> level + let levelH = (ah / 2) >> level + let blocksX8 = (levelW + 7) / 8 + let blocksY8 = (levelH + 7) / 8 + let blocksX32 = (levelW + 31) / 32 + meta[component][level][band] = (count32, blocksX32) + // accumulate_block_mapping's 32x32 count. + count32 += ((blocksX8 + 3) / 4) * ((blocksY8 + 3) / 4) + } + } + } + blockMeta = meta + blockCount32 = count32 + } +} + +/// One parsed frame, CPU side: the per-block payload offset table + the flat payload words the +/// dequant kernel consumes (packet words INCLUDING each 8-byte header, as upstream uploads +/// them), plus the sequence header's facts. +struct ParsedWaveletFrame { + var layout: WaveletLayout + /// Per 32x32 block: u32 word offset into `payload`, or UInt32.max = block missing. + var offsets: [UInt32] + var payload: [UInt32] + var totalBlocks: Int + var decodedBlocks: Int + /// VUI bits from the sequence header (BitstreamSequenceHeader). + var bt2020: Bool + var fullRange: Bool + + /// The frame's Y′CbCr→RGB signal for the presenter's planar CSC. PyroWave today is always + /// BT.709 limited (the host's fixed contract), but the sequence header signals it, so honor + /// what it says. + var cscSignal: CscRows.Signal { + CscRows.Signal(matrix: bt2020 ? 9 : 1, fullRange: fullRange) + } +} + +enum WaveletBitstream { + /// Window kinds of the chunk-aligned framing (host WIN_* constants). + private static let winPacked: UInt16 = 0 + private static let winFragFirst: UInt16 = 1 + private static let winFragCont: UInt16 = 2 + private static let winFragLast: UInt16 = 3 + + /// Parse one AU into the dequant kernel's inputs. `windowSize` > 0 with `chunkAligned` + /// walks the Phase-4 shard-window framing first; otherwise the AU is one packet stream. + /// nil = drop the frame (malformed, no SOF, or not enough blocks survived loss to be worth + /// decoding — upstream's `decoded_blocks > total/2` partial rule). + static func parse(au: Data, chunkAligned: Bool, windowSize: Int) -> ParsedWaveletFrame? { + var state = ParseState() + let ok = au.withUnsafeBytes { (raw: UnsafeRawBufferPointer) -> Bool in + guard let base = raw.baseAddress?.assumingMemoryBound(to: UInt8.self) else { + return false + } + let count = raw.count + if chunkAligned, windowSize >= 8 { + // Whole windows only; a trailing partial window would be a framing bug. + guard count % windowSize == 0 else { return false } + var frag: [UInt8] = [] + var fragLive = false + var pos = 0 + while pos < count { + let win = UnsafeBufferPointer(start: base + pos, count: windowSize) + pos += windowSize + let used = Int(win[0]) | (Int(win[1]) << 8) + let kind = UInt16(win[2]) | (UInt16(win[3]) << 8) + // A zeroed (missing) shard or an overrun drops the window AND breaks any + // fragment chain riding across it (mirrors video_pyrowave.rs push_window). + guard used > 0, 4 + used <= windowSize else { + frag.removeAll(keepingCapacity: true) + fragLive = false + continue + } + let body = UnsafeBufferPointer(start: win.baseAddress! + 4, count: used) + switch kind { + case winPacked: + frag.removeAll(keepingCapacity: true) + fragLive = false + guard state.pushPackets(body) else { return false } + case winFragFirst: + frag.removeAll(keepingCapacity: true) + frag.append(contentsOf: body) + fragLive = true + case winFragCont: + if fragLive { frag.append(contentsOf: body) } + case winFragLast: + if fragLive { + frag.append(contentsOf: body) + let ok = frag.withUnsafeBufferPointer { state.pushPackets($0) } + guard ok else { return false } + } + frag.removeAll(keepingCapacity: true) + fragLive = false + default: + frag.removeAll(keepingCapacity: true) + fragLive = false + } + } + return true + } + return state.pushPackets(UnsafeBufferPointer(start: base, count: count)) + } + guard ok, var frame = state.finish() else { return nil } + // Upstream decode_is_ready(allow_partial=true): with no SOF the frame is undecodable; + // at half the blocks or fewer it is presumed garbage. + guard frame.totalBlocks > 0, frame.decodedBlocks > frame.totalBlocks / 2 else { + return nil + } + // The dequant kernel indexes the offset table by the LAYOUT's block space; the wire's + // total_blocks only counts blocks the encoder emitted. They agree for a full-coverage + // frame, but size the table by the layout. + if frame.offsets.count != frame.layout.blockCount32 { + frame.offsets = Array(frame.offsets.prefix(frame.layout.blockCount32)) + } + return frame + } + + /// Streaming packet-walk state (pyrowave_decoder.cpp push_packet + decode_packet). The + /// SOF sequence header arrives first in every host AU, which fixes the dims → layout → + /// offset-table size before any coefficient packet lands; a coefficient packet before the + /// SOF (its window was lost) is skipped — its block just stays missing. + private struct ParseState { + var layout: WaveletLayout? + var offsets: [UInt32] = [] + var payload: [UInt32] = [] + var totalBlocks = 0 + var decodedBlocks = 0 + var bt2020 = false + var fullRange = false + var sawSOF = false + + mutating func pushPackets(_ buf: UnsafeBufferPointer) -> Bool { + guard let base = buf.baseAddress else { return true } + var pos = 0 + let count = buf.count + while count - pos >= 8 { + let word0 = loadWord(base, pos) + let word1 = loadWord(base, pos + 4) + let extended = (word0 >> 31) & 1 + if extended != 0 { + // BitstreamSequenceHeader: w-1[0:14] h-1[14:28] seq[28:31] ext[31]; + // total[0:24] code[24:26] chroma[26] prim[27] trc[28] mtx[29] range[30] + // siting[31]. + let code = (word1 >> 24) & 0x3 + guard code == 0 else { return false } // only START_OF_FRAME is defined + let chromaRes = (word1 >> 26) & 1 + guard chromaRes == 0 else { return false } // host contract: 4:2:0 + let w = Int(word0 & 0x3fff) + 1 + let h = Int((word0 >> 14) & 0x3fff) + 1 + guard w >= 2, h >= 2, w % 2 == 0, h % 2 == 0 else { return false } + if sawSOF { + // One frame, one geometry — a second SOF must agree. + guard layout?.width == w, layout?.height == h else { return false } + } else { + sawSOF = true + let l = WaveletLayout(width: w, height: h) + layout = l + offsets = [UInt32](repeating: .max, count: l.blockCount32) + payload.reserveCapacity(64 * 1024 / 4) + totalBlocks = Int(word1 & 0xff_ffff) + bt2020 = (word1 >> 29) & 1 != 0 + fullRange = (word1 >> 30) & 1 == 0 // YCBCR_RANGE_FULL = 0 + } + pos += 8 + continue + } + // BitstreamHeader: ballot[0:16] payload_words[16:28] seq[28:31] ext[31]; + // quant_code[0:8] block_index[8:32]. payload_words counts u32s INCLUDING the + // 8-byte header. + let payloadWords = Int((word0 >> 16) & 0xfff) + guard payloadWords >= 2, pos + payloadWords * 4 <= count else { return false } + let blockIndex = Int(word1 >> 8) + if let layout, blockIndex < layout.blockCount32 { + // First write wins (duplicate packets are ignored, like upstream). + if offsets[blockIndex] == .max { + offsets[blockIndex] = UInt32(payload.count) + decodedBlocks += 1 + payload.reserveCapacity(payload.count + payloadWords) + for w in 0.., _ offset: Int) -> UInt32 { + UInt32(base[offset]) + | (UInt32(base[offset + 1]) << 8) + | (UInt32(base[offset + 2]) << 16) + | (UInt32(base[offset + 3]) << 24) + } + + func finish() -> ParsedWaveletFrame? { + guard let layout else { return nil } + return ParsedWaveletFrame( + layout: layout, offsets: offsets, payload: payload, + totalBlocks: totalBlocks, decodedBlocks: decodedBlocks, + bt2020: bt2020, fullRange: fullRange) + } + } +} + +/// One decoded frame's output planes, handed to the presenter's planar render path. The +/// textures belong to the decoder's ring — ring depth (4) plus same-queue hazard tracking keep +/// them valid while referenced. +struct WaveletPlanes { + let y: MTLTexture + let cb: MTLTexture + let cr: MTLTexture + let csc: CscUniform + var width: Int { y.width } + var height: Int { y.height } +} + +final class MetalWaveletDecoder { + /// Matches the Vulkan client's ring: deep enough that a slot is never rewritten while the + /// presenter still samples it in practice; same-queue hazard tracking is the hard backstop. + private static let ringDepth = 4 + + /// Device-capability gate for advertisement (SessionModel) and the settings picker: the + /// dequant kernel needs simdgroup prefix sums with its 16 header lanes inside one + /// simdgroup, so compile the real kernels once and check the pipeline facts. Apple6 (A13) + /// and every Mac2 device pass the family check; the compile probe is authoritative. + static let supported: Bool = { + guard let device = MTLCreateSystemDefaultDevice() else { return false } + guard device.supportsFamily(.apple6) || device.supportsFamily(.mac2) else { return false } + do { + let lib = try device.makeLibrary(source: waveletShaderSource, options: nil) + guard let dequant = lib.makeFunction(name: "wavelet_dequant") else { return false } + let p = try device.makeComputePipelineState(function: dequant) + var shift = false + let fc = MTLFunctionConstantValues() + fc.setConstantValue(&shift, type: .bool, index: 0) + _ = try lib.makeFunction(name: "idwt", constantValues: fc) + return p.threadExecutionWidth >= 16 && p.maxTotalThreadsPerThreadgroup >= 128 + } catch { + waveletLog.info("pyrowave probe: kernels rejected (\(error, privacy: .public))") + return false + } + }() + + private let device: MTLDevice + private let queue: MTLCommandQueue + private let dequantPipeline: MTLComputePipelineState + private let idwtPipeline: MTLComputePipelineState + private let idwtShiftPipeline: MTLComputePipelineState + private let mirrorSampler: MTLSamplerState + + // Size-dependent state, rebuilt when the SOF dims change (this is also the mid-stream + // Reconfigure/resize path — the wavelet decoder is fixed-size per geometry). + private var layout: WaveletLayout? + /// coefficients[component][level]: 4-slice R16Float (levels 0–1) / R32Float (levels 2–4) + /// texture2d_array — the band images (precision-1 split, see MetalWaveletShaders). + private var coefficients: [[MTLTexture]] = [] + /// llViews[component][level]: slice-0 (LL band) 2D write view of `coefficients` — the iDWT + /// output target chaining level L+1 into level L. + private var llViews: [[MTLTexture]] = [] + + private struct Slot { + var y: MTLTexture + var cb: MTLTexture + var cr: MTLTexture + var offsets: MTLBuffer + var payload: MTLBuffer + } + + private var slots: [Slot] = [] + private var nextSlot = 0 + + /// The pump thread owns `decode`; everything mutable is confined to it. + init?(device: MTLDevice, queue: MTLCommandQueue) { + self.device = device + self.queue = queue + do { + let lib = try device.makeLibrary(source: waveletShaderSource, options: nil) + guard let dequantFn = lib.makeFunction(name: "wavelet_dequant") else { return nil } + dequantPipeline = try device.makeComputePipelineState(function: dequantFn) + var shift = false + let fcOff = MTLFunctionConstantValues() + fcOff.setConstantValue(&shift, type: .bool, index: 0) + idwtPipeline = try device.makeComputePipelineState( + function: try lib.makeFunction(name: "idwt", constantValues: fcOff)) + shift = true + let fcOn = MTLFunctionConstantValues() + fcOn.setConstantValue(&shift, type: .bool, index: 0) + idwtShiftPipeline = try device.makeComputePipelineState( + function: try lib.makeFunction(name: "idwt", constantValues: fcOn)) + } catch { + waveletLog.error("pyrowave: pipeline build failed (\(error, privacy: .public))") + return nil + } + guard dequantPipeline.threadExecutionWidth >= 16, + dequantPipeline.maxTotalThreadsPerThreadgroup >= 128 + else { return nil } + // Upstream's mirror_repeat_sampler: mirrored repeat, NEAREST everything, normalized + // coords — the idwt gather footprint + coordinate nudge depend on exactly this. + let samp = MTLSamplerDescriptor() + samp.sAddressMode = .mirrorRepeat + samp.tAddressMode = .mirrorRepeat + samp.minFilter = .nearest + samp.magFilter = .nearest + samp.mipFilter = .notMipmapped + samp.normalizedCoordinates = true + guard let sampler = device.makeSamplerState(descriptor: samp) else { return nil } + mirrorSampler = sampler + } + + /// Decode one AU. Synchronous CPU parse + async GPU decode: returns false when the frame + /// was dropped (malformed / SOF lost / not enough blocks); on true, `completion` fires on a + /// Metal callback thread once the planes are decoded (nil = the GPU pass errored). + /// PUMP THREAD only. + func decode( + au: Data, chunkAligned: Bool, windowSize: Int, + completion: @escaping @Sendable (WaveletPlanes?) -> Void + ) -> Bool { + guard + let frame = WaveletBitstream.parse( + au: au, chunkAligned: chunkAligned, windowSize: windowSize) + else { return false } + + if layout?.width != frame.layout.width || layout?.height != frame.layout.height { + guard rebuild(layout: frame.layout) else { return false } + } + guard let layout, !slots.isEmpty else { return false } + + var slot = slots[nextSlot] + // Grow the payload buffer to the frame (+16-byte zeroed guard: the kernel's 64-bit + // sign-window load and eager plane-byte prefetch may read past the payload end — + // upstream pads its Vulkan buffer for exactly this). + let payloadBytes = frame.payload.count * 4 + if slot.payload.length < payloadBytes + 16 { + guard + let grown = device.makeBuffer( + length: max(64 * 1024, (payloadBytes + 16) * 2), options: .storageModeShared) + else { return false } + slot.payload = grown + slots[nextSlot] = slot + } + frame.offsets.withUnsafeBytes { src in + slot.offsets.contents().copyMemory( + from: src.baseAddress!, byteCount: min(src.count, slot.offsets.length)) + } + frame.payload.withUnsafeBytes { src in + slot.payload.contents().copyMemory(from: src.baseAddress!, byteCount: src.count) + } + memset(slot.payload.contents() + payloadBytes, 0, 16) + + guard let cmd = queue.makeCommandBuffer() else { return false } + + // Stage 1: dequant — every (component, level, band) block grid in one concurrent + // encoder (each dispatch writes its own band layer; no intra-stage hazards, exactly + // like the barrier-free Vulkan dispatch loop). + guard let dequant = cmd.makeComputeCommandEncoder(dispatchType: .concurrent) else { + return false + } + dequant.label = "pyrowave dequant" + dequant.setComputePipelineState(dequantPipeline) + dequant.setBuffer(slot.offsets, offset: 0, index: 0) + dequant.setBuffer(slot.payload, offset: 0, index: 1) + for level in 0...stride, index: 2) + dequant.dispatchThreadgroups( + MTLSize(width: (w + 31) / 32, height: (h + 31) / 32, depth: 1), + threadsPerThreadgroup: MTLSize(width: 128, height: 1, depth: 1)) + } + } + } + dequant.endEncoding() + + // Stage 2: iDWT, coarsest level in — one encoder per level; the encoder boundary is + // the write→sampled-read barrier chaining each level's LL into the next. + for inputLevel in stride(from: WaveletLayout.decompositionLevels - 1, through: 0, by: -1) { + guard let idwt = cmd.makeComputeCommandEncoder(dispatchType: .concurrent) else { + return false + } + idwt.label = "pyrowave idwt L\(inputLevel)" + idwt.setSamplerState(mirrorSampler, index: 0) + // Resolution rides TRANSPOSED (the kernel transposes on load and store). + let rx = layout.levelHeight(inputLevel) + let ry = layout.levelWidth(inputLevel) + var regs = IdwtRegisters( + resolution: SIMD2(Int32(rx), Int32(ry)), + invResolution: SIMD2(1.0 / Float(rx), 1.0 / Float(ry))) + idwt.setBytes(®s, length: MemoryLayout.stride, index: 0) + let grid = MTLSize(width: (rx + 15) / 16, height: (ry + 15) / 16, depth: 1) + let group = MTLSize(width: 64, height: 1, depth: 1) + if inputLevel == 0 { + // 4:2:0: the final full-res pass is luma only (chroma finished at level 1). + idwt.setComputePipelineState(idwtShiftPipeline) + idwt.setTexture(coefficients[0][0], index: 0) + idwt.setTexture(slot.y, index: 1) + idwt.dispatchThreadgroups(grid, threadsPerThreadgroup: group) + } else { + for component in 0..<3 { + idwt.setTexture(coefficients[component][inputLevel], index: 0) + if component != 0 && inputLevel == 1 { + // 4:2:0 chroma emits its final half-res plane one level early. + idwt.setComputePipelineState(idwtShiftPipeline) + idwt.setTexture(component == 1 ? slot.cb : slot.cr, index: 1) + } else { + idwt.setComputePipelineState(idwtPipeline) + idwt.setTexture(llViews[component][inputLevel - 1], index: 1) + } + idwt.dispatchThreadgroups(grid, threadsPerThreadgroup: group) + } + } + idwt.endEncoding() + } + + let planes = WaveletPlanes( + y: slot.y, cb: slot.cb, cr: slot.cr, + csc: CscRows.rows(frame.cscSignal, depth: 8, msbPacked: false)) + cmd.addCompletedHandler { buffer in + completion(buffer.error == nil ? planes : nil) + } + cmd.commit() + nextSlot = (nextSlot + 1) % Self.ringDepth + return true + } + + /// (Re)allocate every size-dependent resource for `layout`'s geometry. Also the mid-stream + /// resize path: a Reconfigure shows up here as new SOF dims. + private func rebuild(layout newLayout: WaveletLayout) -> Bool { + waveletLog.info( + "pyrowave: building decoder \(newLayout.width)x\(newLayout.height) (aligned \(newLayout.alignedWidth)x\(newLayout.alignedHeight), \(newLayout.blockCount32) blocks)") + var coeff: [[MTLTexture]] = [] + var lls: [[MTLTexture]] = [] + for component in 0..<3 { + var perLevel: [MTLTexture] = [] + var perLevelLL: [MTLTexture] = [] + for level in 0.. MTLTexture? in + let desc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .r8Unorm, width: w, height: h, mipmapped: false) + desc.usage = [.shaderRead, .shaderWrite] + desc.storageMode = .private + let t = self.device.makeTexture(descriptor: desc) + t?.label = name + return t + } + guard + let y = plane(newLayout.width, newLayout.height, "pyrowave Y[\(i)]"), + let cb = plane(newLayout.width / 2, newLayout.height / 2, "pyrowave Cb[\(i)]"), + let cr = plane(newLayout.width / 2, newLayout.height / 2, "pyrowave Cr[\(i)]"), + let offsets = device.makeBuffer( + length: max(newLayout.blockCount32 * 4, 4), options: .storageModeShared), + let payload = device.makeBuffer(length: 64 * 1024, options: .storageModeShared) + else { return false } + newSlots.append(Slot(y: y, cb: cb, cr: cr, offsets: offsets, payload: payload)) + } + + coefficients = coeff + llViews = lls + slots = newSlots + nextSlot = 0 + layout = newLayout + return true + } + + // MSL-side layouts (MetalWaveletShaders.swift) — keep in lockstep. + private struct DequantRegisters { + var resolution: SIMD2 + var outputLayer: Int32 + var blockOffset32x32: Int32 + var blockStride32x32: Int32 + } + + private struct IdwtRegisters { + var resolution: SIMD2 + var invResolution: SIMD2 + } +} +#endif diff --git a/clients/apple/Sources/PunktfunkKit/Video/MetalWaveletShaders.swift b/clients/apple/Sources/PunktfunkKit/Video/MetalWaveletShaders.swift new file mode 100644 index 00000000..78cbc39c --- /dev/null +++ b/clients/apple/Sources/PunktfunkKit/Video/MetalWaveletShaders.swift @@ -0,0 +1,551 @@ +// PyroWave decode compute kernels — the Metal port of the vendored Vulkan shaders +// (crates/pyrowave-sys/vendor/pyrowave/shaders/wavelet_dequant.comp + idwt.comp, upstream pin +// 509e4f88, MIT © 2025 Hans-Kristian Arntzen). Runtime-compiled Swift strings per client +// convention (no metallib build step — see GamepadChrome.swift's rationale); these are the +// client's first compute pipelines. +// +// Port notes (design/pyrowave-codec-plan.md §4.7): +// • Only the STORAGE_MODE 0 path exists: MSL device pointers replace the 8/16-bit-storage SSBO +// aliases; the texel-buffer (mode 1) and linear-image (mode 2) fallbacks are non-Apple IHV +// workarounds and are dropped, as is the fragment-iDWT path (Mali/Adreno only). +// • Subgroup ops map 1:1: subgroupInclusiveAdd → simd_prefix_inclusive_sum, and the fixed +// 32-wide Apple simdgroups take the GLSL's `SubgroupSize <= 32` scan branch; the shuffle-up +// and LDS fallbacks for exotic wave sizes are dead code here. The dequant kernel needs the +// 16 header lanes inside ONE simdgroup — MetalWaveletDecoder's probe enforces +// threadExecutionWidth >= 16. +// • Precision matches upstream's desktop default (PYROWAVE_PRECISION=1): float arithmetic, +// half2 threadgroup storage; the coefficient textures are R16Float for DWT levels 0–1 and +// R32Float for levels 2–4 (the low-res levels feed long reconstruction chains — upstream +// keeps them fp32 for exactly that reason). +// • The gather + mirrored-repeat addressing in idwt is the precision-sensitive spot (upstream +// fought a Mali compiler bug there); the golden-frame PSNR fixtures are the guard. + +import Foundation + +let waveletShaderSource = """ +#include +using namespace metal; + +// --------------------------------------------------------------------------------------------- +// Shared helpers (dwt_swizzle.h / constants.h / dwt_quant_scale.h) +// --------------------------------------------------------------------------------------------- + +static inline int2 unswizzle8x8(uint index) +{ + uint y = extract_bits(index, 0, 1); + uint x = extract_bits(index, 1, 2); + y |= extract_bits(index, 3, 2) << 1; + x |= extract_bits(index, 5, 1) << 2; + return int2(int(x), int(y)); +} + +// GLSL bitfieldExtract(x, 0, n) where n may be 0; MSL extract_bits(bits=0) is not guaranteed +// to return 0, so mask explicitly. +static inline uint mask_lo(uint x, int n) +{ + return (n <= 0) ? 0u : (x & (0xffffffffu >> (32 - n))); +} + +// pyrowave_common.hpp decode_quant: custom FP formulation, MaxScaleExp = 4. +static inline float decode_quant(uint quant_code) +{ + int e = 4 - int(quant_code >> 3); + int m = int(quant_code) & 0x7; + return (1.0f / (8.0f * 1024.0f * 1024.0f)) * float((8 + m) * (1 << (20 + e))); +} + +// dwt_quant_scale.h: per-8x8 quant scale, min 0.25, max ~2.2. +static inline float decode_quant_scale(uint code) +{ + return float(code) / 8.0f + 0.25f; +} + +// constants.h +constant int QUANT_SCALE_OFFSET = 20; +constant int QUANT_SCALE_BITS = 4; + +// --------------------------------------------------------------------------------------------- +// wavelet_dequant — one 128-thread threadgroup decodes one 32x32 coefficient block +// --------------------------------------------------------------------------------------------- + +struct DequantRegisters { + int2 resolution; + int output_layer; + int block_offset_32x32; + int block_stride_32x32; +}; + +struct DecodedPair { float4 col0; float4 col1; }; // GLSL mat2x4: m[j][i] -> colJ[i] + +// Bit-plane magnitude decode for one thread's 4x2 coefficient group (decode_payload in the +// GLSL). `code_word` is the 8x8 block's 16-bit control word (2 bits of extra planes per 4x2 +// group), `q_bits` the base plane count, `offset` the block's plane-payload start byte, +// `block_index` this thread's group (0..7). Nonzero magnitudes get the +0.5 deadzone +// reconstruction bias. +static DecodedPair decode_payload(const device uchar *payload_u8, + uint code_word, uint q_bits, uint offset, uint block_index) +{ + DecodedPair m; + m.col0 = float4(0.0f); + m.col1 = float4(0.0f); + if (code_word == 0) + return m; + + int bit_offset = 2 * int(block_index); + + uint lsbs = code_word & 0x5555u; + uint msbs = code_word & 0xaaaau; + uint msbs_shift = msbs >> 1; + msbs |= msbs_shift; + + uint byte_offset = + popcount(mask_lo(lsbs, bit_offset)) + + popcount(mask_lo(msbs, bit_offset)) + + q_bits * block_index + offset; + + uint payload = uint(payload_u8[byte_offset]); + + uint local_control_word = extract_bits(code_word, uint(bit_offset), 2); + int decoded_abs[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + int plane_iterations = int(q_bits + local_control_word); + + for (int q = plane_iterations - 1; q >= 0; q--) + { + for (int b = 0; b < 8; b++) + { + int decoded = int(extract_bits(payload, uint(b), 1)); + decoded_abs[b] = insert_bits(decoded_abs[b], decoded, uint(q), 1); + } + byte_offset++; + payload = uint(payload_u8[byte_offset]); + } + + for (int i = 0; i < 4; i++) + { + for (int j = 0; j < 2; j++) + { + float v = float(decoded_abs[i * 2 + j]); + if (v != 0.0f) + v += 0.5f; + if (j == 0) m.col0[i] = v; else m.col1[i] = v; + } + } + return m; +} + +kernel void wavelet_dequant( + texture2d_array uDequantImg [[texture(0)]], + const device uint *payload_offsets [[buffer(0)]], + const device uint *payload_u32 [[buffer(1)]], + constant DequantRegisters ®isters [[buffer(2)]], + uint3 wg_id [[threadgroup_position_in_grid]], + uint local_index [[thread_index_in_threadgroup]], + uint simd_lane [[thread_index_in_simdgroup]], + uint simd_group [[simdgroup_index_in_threadgroup]], + uint simd_size [[threads_per_simdgroup]]) +{ + // STORAGE_MODE 0's three aliased SSBO views over one buffer, as typed pointers. + const device ushort *payload_u16 = reinterpret_cast(payload_u32); + const device uchar *payload_u8 = reinterpret_cast(payload_u32); + + threadgroup uint shared_sign_offset; + threadgroup uint shared_plane_byte_offsets[16]; + threadgroup uint shared_sign_scan[128 / 4]; + + int block_index_32x32 = int(uint(registers.block_offset_32x32) + + wg_id.y * uint(registers.block_stride_32x32) + + wg_id.x); + + uint block_local_index = extract_bits(local_index, 0, 3); + uint block_x = extract_bits(local_index, 3, 2); + uint block_y = extract_bits(local_index, 5, 2); + uint linear_block = block_y * 4 + block_x; + + // Each thread individually decodes 8 values (a 4x2 group of its 8x8 block). + int2 local_coord = unswizzle8x8(block_local_index << 3); + + int2 coord = int2(wg_id.xy) * 32; + coord += 8 * int2(int(block_x), int(block_y)); + coord += local_coord; + + uint offset_u32 = payload_offsets[block_index_32x32]; + + // Missing / lost block: zero coefficients (this is how a partial frame's holes decode). + if (offset_u32 == ~0u) + { + for (int j = 0; j < 2; j++) + for (int i = 0; i < 4; i++) + uDequantImg.write(float4(0.0f), uint2(coord + int2(i, j)), uint(registers.output_layer)); + return; + } + + uint ballot = payload_u32[offset_u32] & 0xffffu; + uint q_code = payload_u32[offset_u32 + 1] & 0xffu; + + // Threads 0..15 (one per 8x8 block, all inside simdgroup 0) prefix-scan the per-block + // plane-payload byte costs into shared_plane_byte_offsets, and lane 15 records where the + // sign bitstream starts. + if (local_index < 16) + { + uint control_word = 0; + uint q_bits = 0; + + if (extract_bits(ballot, local_index, 1) != 0) + { + uint local_code_offset = popcount(mask_lo(ballot, int(local_index))); + control_word = uint(payload_u16[offset_u32 * 2 + 4 + local_code_offset]); + q_bits = uint(payload_u8[offset_u32 * 4 + 8 + popcount(ballot) * 2 + local_code_offset]) & 0xfu; + } + + uint lsbs = control_word & 0x5555u; + uint msbs = control_word & 0xaaaau; + uint msbs_shift = msbs >> 1; + msbs |= msbs_shift; + uint byte_cost = popcount(lsbs) + popcount(msbs) + q_bits * 8; + + uint byte_scan = offset_u32 * 4 + 8 + 3 * popcount(ballot) + simd_prefix_inclusive_sum(byte_cost); + if (local_index == 15) + shared_sign_offset = 8 * byte_scan; + shared_plane_byte_offsets[local_index] = byte_scan - byte_cost; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + DecodedPair v; + int significant_count; + + if (extract_bits(ballot, linear_block, 1) != 0) + { + uint local_code_offset = popcount(mask_lo(ballot, int(linear_block))); + + uint control_word = uint(payload_u16[offset_u32 * 2 + 4 + local_code_offset]); + uint control_word2 = uint(payload_u8[offset_u32 * 4 + 8 + popcount(ballot) * 2 + local_code_offset]); + + v = decode_payload(payload_u8, control_word, control_word2 & 0xfu, + shared_plane_byte_offsets[linear_block], block_local_index); + + significant_count = 0; + for (int j = 0; j < 2; j++) + for (int i = 0; i < 4; i++) + significant_count += int(((j == 0) ? v.col0[i] : v.col1[i]) != 0.0f); + + float q = decode_quant(q_code); + float inv_scale = q * decode_quant_scale(extract_bits(control_word2, uint(QUANT_SCALE_OFFSET - 16), uint(QUANT_SCALE_BITS))); + + v.col0 *= inv_scale; + v.col1 *= inv_scale; + } + else + { + v.col0 = float4(0.0f); + v.col1 = float4(0.0f); + significant_count = 0; + } + + // Cross-threadgroup scan of significant-coefficient counts → each thread's first sign-bit + // position. Apple simdgroups are >= 16 wide, so this is the GLSL's `SubgroupSize <= 32` + // branch; the shuffle/LDS fallbacks are unnecessary. + int significant_scan = int(simd_prefix_inclusive_sum(uint(significant_count))); + if (simd_lane == simd_size - 1) + shared_sign_scan[simd_group] = uint(significant_scan); + + threadgroup_barrier(mem_flags::mem_threadgroup); + uint num_simdgroups = (128 + simd_size - 1) / simd_size; + if (local_index < num_simdgroups) + shared_sign_scan[local_index] = simd_prefix_inclusive_sum(shared_sign_scan[local_index]); + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint sign_offset = shared_sign_offset + uint(significant_scan - significant_count); + if (simd_group != 0) + sign_offset += shared_sign_scan[simd_group - 1]; + + // Load 64 bits of sign stream and bit-align (may read one word past the payload — the + // buffer carries a 16-byte zeroed guard tail for exactly this). + uint sign_word = payload_u32[sign_offset / 32 + 0]; + uint sign_word_upper = payload_u32[sign_offset / 32 + 1]; + + uint masked_sign_offset = sign_offset & 31u; + if (masked_sign_offset != 0) + { + sign_word >>= masked_sign_offset; + sign_word |= sign_word_upper << (32 - masked_sign_offset); + } + + int sign_counter = 0; + + for (int i = 0; i < 4; i++) + { + for (int j = 0; j < 2; j++) + { + float val = (j == 0) ? v.col0[i] : v.col1[i]; + if (val != 0.0f) + { + val *= 1.0f - 2.0f * float(extract_bits(sign_word, uint(sign_counter), 1)); + sign_counter++; + if (j == 0) v.col0[i] = val; else v.col1[i] = val; + } + } + } + + for (int j = 0; j < 2; j++) + for (int i = 0; i < 4; i++) + uDequantImg.write(float4((j == 0) ? v.col0[i] : v.col1[i]), + uint2(coord + int2(i, j)), uint(registers.output_layer)); +} + +// --------------------------------------------------------------------------------------------- +// idwt — inverse CDF 9/7; one 64-thread threadgroup reconstructs one 32x32 output tile from the +// four half-res band layers (LL/HL/LH/HH), with a 4-sample mirror apron. The caller passes the +// band-image resolution TRANSPOSED (the kernel transposes on load and store, so one kernel does +// both the horizontal and vertical passes). +// --------------------------------------------------------------------------------------------- + +constant bool DCShift [[function_constant(0)]]; + +struct IdwtRegisters { + int2 resolution; + float2 inv_resolution; +}; + +constant int APRON = 4; +constant int APRON_HALF = APRON / 2; +constant int BLOCK_SIZE = 32; +constant int BLOCK_SIZE_HALF = BLOCK_SIZE >> 1; + +// CDF 9/7 lifting constants (dwt_common.h). +constant float ALPHA = -1.586134342059924f; +constant float BETA = -0.052980118572961f; +constant float GAMMA = 0.882911075530934f; +constant float DELTA = 0.443506852043971f; +constant float K = 1.230174104914001f; +constant float inv_K = 1.0f / 1.230174104914001f; + +constant int SHARED_ROWS = (BLOCK_SIZE + 2 * APRON) / 2; // 20 +constant int SHARED_COLS = (BLOCK_SIZE + 2 * APRON) + 1; // 41 (+1 avoids bank conflicts) + +static inline float2 load_shared(threadgroup half2 (&blk)[SHARED_ROWS][SHARED_COLS], int y, int x) +{ + return float2(blk[y][x]); +} + +static inline void store_shared(threadgroup half2 (&blk)[SHARED_ROWS][SHARED_COLS], int y, int x, float2 v) +{ + blk[y][x] = half2(v); +} + +// Even/odd-phase coordinate nudge so mirrored-repeat gather reproduces JPEG2000 whole-sample +// mirroring at the image borders, then transpose (uv.yx) on load. +static inline float2 generate_mirror_uv(int2 coord, bool even_x, bool even_y, + int2 resolution, float2 inv_resolution) +{ + coord.x -= int(even_x && coord.x < 0); + coord.y -= int(even_y && coord.y < 0); + coord += 1; + coord.x += int(!even_x && coord.x >= resolution.x); + coord.y += int(!even_y && coord.y >= resolution.y); + float2 uv = float2(coord) * inv_resolution; + return uv.yx; +} + +static inline void write_shared_4x4(threadgroup half2 (&blk)[SHARED_ROWS][SHARED_COLS], + int2 coord, float4 t0, float4 t1, float4 t2, float4 t3) +{ + store_shared(blk, coord.y + 0, 2 * coord.x + 0, float2(t0.x, t2.x)); + store_shared(blk, coord.y + 0, 2 * coord.x + 1, float2(t1.x, t3.x)); + store_shared(blk, coord.y + 0, 2 * coord.x + 2, float2(t0.y, t2.y)); + store_shared(blk, coord.y + 0, 2 * coord.x + 3, float2(t1.y, t3.y)); + store_shared(blk, coord.y + 1, 2 * coord.x + 0, float2(t0.z, t2.z)); + store_shared(blk, coord.y + 1, 2 * coord.x + 1, float2(t1.z, t3.z)); + store_shared(blk, coord.y + 1, 2 * coord.x + 2, float2(t0.w, t2.w)); + store_shared(blk, coord.y + 1, 2 * coord.x + 3, float2(t1.w, t3.w)); +} + +// textureGather(...).wxzy — Metal's gather returns the same counter-clockwise-from-(i0,j1) +// component order as Vulkan, so the reorder is identical. +static inline float4 gather_layer(texture2d_array tex, sampler smp, + float2 uv, uint layer) +{ + float4 g = tex.gather(smp, uv, layer); + return float4(g.w, g.x, g.z, g.y); +} + +static void load_image_with_apron(texture2d_array tex, sampler smp, + threadgroup half2 (&blk)[SHARED_ROWS][SHARED_COLS], + uint local_index, uint2 wg_id, + int2 resolution, float2 inv_resolution) +{ + int2 base_coord = int2(wg_id) * BLOCK_SIZE_HALF - APRON_HALF; + int2 local_coord0 = 2 * unswizzle8x8(local_index); + int2 coord0 = base_coord + local_coord0; + + // Band layers gathered in 0/2/1/3 order (LL/LH/HL/HH interleave for the 2x2 scatter). + float4 texels0 = gather_layer(tex, smp, generate_mirror_uv(coord0, true, true, resolution, inv_resolution), 0); + float4 texels1 = gather_layer(tex, smp, generate_mirror_uv(coord0, false, true, resolution, inv_resolution), 2); + float4 texels2 = gather_layer(tex, smp, generate_mirror_uv(coord0, true, false, resolution, inv_resolution), 1); + float4 texels3 = gather_layer(tex, smp, generate_mirror_uv(coord0, false, false, resolution, inv_resolution), 3); + write_shared_4x4(blk, local_coord0, texels0, texels1, texels2, texels3); + + int2 local_coord_horiz = int2(BLOCK_SIZE_HALF + 2 * int(local_index % 2u), 2 * int(local_index / 2u)); + if (local_coord_horiz.y < BLOCK_SIZE_HALF + 2 * APRON_HALF) + { + int2 c = base_coord + local_coord_horiz; + texels0 = gather_layer(tex, smp, generate_mirror_uv(c, true, true, resolution, inv_resolution), 0); + texels1 = gather_layer(tex, smp, generate_mirror_uv(c, false, true, resolution, inv_resolution), 2); + texels2 = gather_layer(tex, smp, generate_mirror_uv(c, true, false, resolution, inv_resolution), 1); + texels3 = gather_layer(tex, smp, generate_mirror_uv(c, false, false, resolution, inv_resolution), 3); + write_shared_4x4(blk, local_coord_horiz, texels0, texels1, texels2, texels3); + } + + int2 local_coord_vert = local_coord_horiz.yx; + if (local_coord_vert.x < BLOCK_SIZE_HALF) + { + int2 c = base_coord + local_coord_vert; + texels0 = gather_layer(tex, smp, generate_mirror_uv(c, true, true, resolution, inv_resolution), 0); + texels1 = gather_layer(tex, smp, generate_mirror_uv(c, false, true, resolution, inv_resolution), 2); + texels2 = gather_layer(tex, smp, generate_mirror_uv(c, true, false, resolution, inv_resolution), 1); + texels3 = gather_layer(tex, smp, generate_mirror_uv(c, false, false, resolution, inv_resolution), 3); + write_shared_4x4(blk, local_coord_vert, texels0, texels1, texels2, texels3); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); +} + +static void inverse_transform8x2(threadgroup half2 (&blk)[SHARED_ROWS][SHARED_COLS], uint local_index) +{ + const int SIZE = 8; + const int PADDED_SIZE = SIZE + 2 * APRON; + const int PADDED_SIZE_HALF = PADDED_SIZE / 2; + float2 values[PADDED_SIZE]; + + int2 local_coord = int2(8 * int(local_index % 4u), int(local_index / 4u)); + + for (int i = 0; i < PADDED_SIZE; i += 2) + { + float2 v0 = load_shared(blk, local_coord.y, local_coord.x + i + 0); + float2 v1 = load_shared(blk, local_coord.y, local_coord.x + i + 1); + values[i + 0] = v0 * K; + values[i + 1] = v1 * inv_K; + } + + // CDF 9/7 inverse lifting steps. + for (int i = 2; i < PADDED_SIZE - 1; i += 2) + values[i] -= DELTA * (values[i - 1] + values[i + 1]); + for (int i = 3; i < PADDED_SIZE - 2; i += 2) + values[i] -= GAMMA * (values[i - 1] + values[i + 1]); + for (int i = 4; i < PADDED_SIZE - 3; i += 2) + values[i] -= BETA * (values[i - 1] + values[i + 1]); + for (int i = 5; i < PADDED_SIZE - 4; i += 2) + values[i] -= ALPHA * (values[i - 1] + values[i + 1]); + + // Avoid WAR hazard. + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++) + { + float2 a = values[2 * i + 0]; + float2 b = values[2 * i + 1]; + + // Transpose the 2x2 block, transpose write. + float2 t0 = float2(a.x, b.x); + float2 t1 = float2(a.y, b.y); + + int y_coord = (local_coord.x >> 1) + (i - APRON_HALF); + store_shared(blk, y_coord, 2 * local_coord.y + 0, t0); + store_shared(blk, y_coord, 2 * local_coord.y + 1, t1); + } +} + +static void inverse_transform4x2(threadgroup half2 (&blk)[SHARED_ROWS][SHARED_COLS], + uint local_index, bool active_lane, int y_offset) +{ + const int SIZE = 4; + const int PADDED_SIZE = SIZE + 2 * APRON; + const int PADDED_SIZE_HALF = PADDED_SIZE / 2; + float2 values[PADDED_SIZE]; + + int2 local_coord = int2(4 * int(local_index % 8u), int(local_index / 8u) + y_offset); + + if (active_lane) + { + for (int i = 0; i < PADDED_SIZE; i += 2) + { + float2 v0 = load_shared(blk, local_coord.y, local_coord.x + i + 0); + float2 v1 = load_shared(blk, local_coord.y, local_coord.x + i + 1); + values[i + 0] = v0 * K; + values[i + 1] = v1 * inv_K; + } + + for (int i = 2; i < PADDED_SIZE - 1; i += 2) + values[i] -= DELTA * (values[i - 1] + values[i + 1]); + for (int i = 3; i < PADDED_SIZE - 2; i += 2) + values[i] -= GAMMA * (values[i - 1] + values[i + 1]); + for (int i = 4; i < PADDED_SIZE - 3; i += 2) + values[i] -= BETA * (values[i - 1] + values[i + 1]); + for (int i = 5; i < PADDED_SIZE - 4; i += 2) + values[i] -= ALPHA * (values[i - 1] + values[i + 1]); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (active_lane) + { + for (int i = APRON_HALF; i < PADDED_SIZE_HALF - APRON_HALF; i++) + { + float2 a = values[2 * i + 0]; + float2 b = values[2 * i + 1]; + + float2 t0 = float2(a.x, b.x); + float2 t1 = float2(a.y, b.y); + + int y_coord = (local_coord.x >> 1) + (i - APRON_HALF); + store_shared(blk, y_coord, 2 * local_coord.y + 0, t0); + store_shared(blk, y_coord, 2 * local_coord.y + 1, t1); + } + } +} + +kernel void idwt( + texture2d_array uTexture [[texture(0)]], + texture2d uOutput [[texture(1)]], + sampler uSampler [[sampler(0)]], + constant IdwtRegisters ®isters [[buffer(0)]], + uint3 wg_id [[threadgroup_position_in_grid]], + uint local_index [[thread_index_in_threadgroup]]) +{ + threadgroup half2 shared_block[SHARED_ROWS][SHARED_COLS]; + + load_image_with_apron(uTexture, uSampler, shared_block, local_index, wg_id.xy, + registers.resolution, registers.inv_resolution); + + // Horizontal transform. + inverse_transform8x2(shared_block, local_index); + + // Also need to transform the apron. + inverse_transform4x2(shared_block, local_index, local_index < 32, BLOCK_SIZE_HALF); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Vertical transform. + inverse_transform8x2(shared_block, local_index); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + int2 local_coord = unswizzle8x8(local_index); + + for (int y = local_coord.y; y < BLOCK_SIZE_HALF; y += 8) + { + for (int x = local_coord.x; x < BLOCK_SIZE; x += 8) + { + float2 v = load_shared(shared_block, y, x); + if (DCShift) + v += 0.5f; + // Transposed store (wg_id.yx) — undoes the transpose-on-load; out-of-range writes + // at the aligned-size overhang are dropped by Metal (matching the Vulkan behavior). + int2 out0 = int2(2 * y + 0, x) + BLOCK_SIZE * int2(int(wg_id.y), int(wg_id.x)); + int2 out1 = int2(2 * y + 1, x) + BLOCK_SIZE * int2(int(wg_id.y), int(wg_id.x)); + uOutput.write(float4(v.x), uint2(out0)); + uOutput.write(float4(v.y), uint2(out1)); + } + } +} +""" diff --git a/crates/pf-client-core/src/session.rs b/crates/pf-client-core/src/session.rs index 5cb000fb..d6036e69 100644 --- a/crates/pf-client-core/src/session.rs +++ b/crates/pf-client-core/src/session.rs @@ -328,6 +328,20 @@ fn pump( // Live host↔client clock offset: loaded per frame (Relaxed) so mid-stream re-syncs (an NTP // step, drift) keep the capture-clock latency stats honest — never cached at session start. let clock_offset_live = connector.clock_offset_shared(); + // PUNKTFUNK_DEBUG_RECONFIGURE=WxH@HZ:SECS — lab lever: request ONE mid-stream mode + // switch N seconds in, so a headless session (no window manager to drag a window in) + // can exercise the resize path deterministically — host pipeline rebuild, decoder + // follow-through (e.g. the PyroWave in-place rebuild), overlay/aspect handling. + let pump_start = Instant::now(); + let mut debug_reconfig = std::env::var("PUNKTFUNK_DEBUG_RECONFIGURE") + .ok() + .and_then(|s| { + let parsed = parse_debug_reconfigure(&s); + if parsed.is_none() { + tracing::warn!(value = %s, "PUNKTFUNK_DEBUG_RECONFIGURE not understood (want WxH@HZ:SECS) — ignored"); + } + parsed + }); let mut total_frames = 0u64; // Newest frame index handed to the decoder — the staleness bar for late partials. let mut newest_decoded_idx: Option = None; @@ -368,6 +382,18 @@ fn pump( if stop.load(Ordering::SeqCst) { break None; } + if let Some((mode, delay)) = debug_reconfig { + if pump_start.elapsed() >= delay { + tracing::info!( + ?mode, + "PUNKTFUNK_DEBUG_RECONFIGURE: requesting mid-stream mode switch" + ); + if let Err(e) = connector.request_mode(mode) { + tracing::warn!(error = ?e, "debug mode switch request failed"); + } + debug_reconfig = None; + } + } // 20 ms wait: audio has its own thread now, so this only bounds stop-flag // responsiveness and the per-iteration keyframe-recovery check (a frame arrives // every ~8–16 ms at 60–120 Hz anyway, so this rarely times out mid-stream). @@ -765,3 +791,43 @@ fn spawn_audio( .map_err(|e| tracing::warn!(error = %e, "audio thread failed to start — audio disabled")) .ok() } + +/// Parse the `PUNKTFUNK_DEBUG_RECONFIGURE` lab lever: `WxH@HZ:SECS` → request that mode +/// SECS seconds into the stream (e.g. `1280x720@60:5`). +fn parse_debug_reconfigure(s: &str) -> Option<(Mode, Duration)> { + let (mode_s, secs_s) = s.split_once(':')?; + let (res, hz) = mode_s.split_once('@')?; + let (w, h) = res.split_once('x')?; + let mode = Mode { + width: w.trim().parse().ok()?, + height: h.trim().parse().ok()?, + refresh_hz: hz.trim().parse().ok()?, + }; + Some((mode, Duration::from_secs(secs_s.trim().parse().ok()?))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn debug_reconfigure_parses_the_documented_shape() { + let (mode, delay) = parse_debug_reconfigure("1280x720@60:5").unwrap(); + assert_eq!((mode.width, mode.height, mode.refresh_hz), (1280, 720, 60)); + assert_eq!(delay, Duration::from_secs(5)); + } + + #[test] + fn debug_reconfigure_rejects_garbage() { + for bad in [ + "", + "1280x720", + "1280x720@60", + "x@:", + "ax b@c:d", + "1280x720@60:x", + ] { + assert!(parse_debug_reconfigure(bad).is_none(), "{bad:?} parsed"); + } + } +} diff --git a/crates/pf-client-core/src/video.rs b/crates/pf-client-core/src/video.rs index 0866cc18..07e2cb76 100644 --- a/crates/pf-client-core/src/video.rs +++ b/crates/pf-client-core/src/video.rs @@ -628,6 +628,11 @@ impl Decoder { pub fn decode_frame( &mut self, au: &[u8], + // Only the PyroWave backend reads the flags; without that feature the param is unused. + #[cfg_attr( + not(all(target_os = "linux", feature = "pyrowave")), + allow(unused_variables) + )] user_flags: u32, complete: bool, ) -> Result> { diff --git a/crates/pf-client-core/src/video_pyrowave.rs b/crates/pf-client-core/src/video_pyrowave.rs index 74cce622..f8397981 100644 --- a/crates/pf-client-core/src/video_pyrowave.rs +++ b/crates/pf-client-core/src/video_pyrowave.rs @@ -19,6 +19,14 @@ //! content-equivalent ones from [`VulkanDecodeDevice`]'s exported extension lists, //! feature facts and queue-family shape (pyrowave reads them for extension/feature //! detection; pointer identity is not required). +//! +//! **Mid-stream resize:** the pyrowave decoder object is fixed-size, but every frame's +//! bitstream opens with a sequence header carrying its dimensions — [`au_dims`] sniffs +//! it and [`PyroWaveDecoder::reconfigure`] rebuilds the decoder + plane ring in place +//! when the host's `Reconfigure` pipeline rebuild lands (the pyrowave *device*, command +//! pool and pinned create-infos are dimension-independent and survive). Superseded plane +//! rings are retired, not destroyed — the presenter may still hold their views (see +//! [`RETIRE_HANDOVERS`]). use crate::video::{ColorDesc, VulkanDecodeDevice}; use anyhow::{bail, Context as _, Result}; @@ -27,6 +35,7 @@ use ash::vk::Handle as _; use pyrowave_sys as pw; use std::ffi::{c_char, c_void, CString}; use std::sync::Arc; +use std::time::{Duration, Instant}; /// Plane-set ring depth: decode writes slot N while the presenter may still sample /// N-1/N-2 (its own submission raced ahead under the shared queue's FIFO order, so @@ -34,6 +43,18 @@ use std::sync::Arc; /// keeps LOGICAL reuse far enough behind). const RING: usize = 4; +/// A mid-stream resize retires the old plane ring, but its images can't be destroyed +/// immediately: the pump→presenter frame channel (depth 2, newest-wins) may still hold a +/// frame referencing them, and the presenter binds a frame's views into its descriptor +/// set only inside the `present` call that carries it. Once this many NEW-ring frames +/// have been handed over, every old-ring frame has been displaced from the channel and +/// any present that picked one up has long finished recording; combined with the +/// queue-idle taken before destruction (covers submitted GPU work) the retired images +/// are provably unreachable. The wall-clock floor is a belt for a presenter stalled +/// mid-`present` (swapchain acquire on an occluded window) while frames keep flowing. +const RETIRE_HANDOVERS: u32 = 8; +const RETIRE_MIN_AGE: Duration = Duration::from_millis(250); + fn pw_check(r: pw::pyrowave_result, what: &str) -> Result<()> { if r == pw::pyrowave_result_PYROWAVE_SUCCESS { Ok(()) @@ -42,6 +63,54 @@ fn pw_check(r: pw::pyrowave_result, what: &str) -> Result<()> { } } +/// Parse an upstream `BitstreamSequenceHeader` (pyrowave_common.hpp) at the start of +/// `bytes`: 8 bytes, two LE u32s — word 0 = `width_minus_1:14 | height_minus_1:14 | +/// sequence:3 | extended:1`, word 1 = `total_blocks:24 | code:2 | …`. Returns the frame +/// dimensions when this really is a START-OF-FRAME sequence header (the `extended` bit +/// distinguishes it from a regular `BitstreamHeader`, which carries a wavelet block). +fn seq_header_dims(bytes: &[u8]) -> Option<(u32, u32)> { + if bytes.len() < 8 { + return None; + } + let w0 = u32::from_le_bytes(bytes[0..4].try_into().unwrap()); + let w1 = u32::from_le_bytes(bytes[4..8].try_into().unwrap()); + if w0 >> 31 == 0 { + return None; // regular block header, not a sequence header + } + if (w1 >> 24) & 0x3 != 0 { + return None; // extended, but not BITSTREAM_EXTENDED_CODE_START_OF_FRAME + } + Some(((w0 & 0x3FFF) + 1, ((w0 >> 14) & 0x3FFF) + 1)) +} + +/// The frame dimensions an AU announces, or `None` when they can't be known from this AU +/// (the sequence header rode a lost shard of a partial). The encoder writes exactly one +/// sequence header per frame, at byte 0 of the frame's bitstream — so it sits at the +/// start of an unaligned AU, and at the start of the FIRST window's body in a +/// chunk-aligned AU (§4.4 framing: 4-byte prefix `used:u16 | kind:u16`; kind PACKED or +/// FRAG_FIRST both begin with the frame's first packet, and that packet begins with the +/// sequence header). +fn au_dims(au: &[u8], aligned: bool, wire_window: usize) -> Option<(u32, u32)> { + if !aligned { + return seq_header_dims(au); + } + let win = &au[..au.len().min(wire_window)]; + if win.len() < 4 { + return None; + } + let used = u16::from_le_bytes([win[0], win[1]]) as usize; + let kind = u16::from_le_bytes([win[2], win[3]]); + if used == 0 || 4 + used > win.len() { + return None; // first window lost/garbage — the sequence header went with it + } + // WIN_PACKED (0) and WIN_FRAG_FIRST (1) both start at the frame's first packet; + // a CONT/LAST fragment here would mean the first window was lost. + if kind > 1 { + return None; + } + seq_header_dims(&win[4..4 + used]) +} + /// Content-equivalent reconstruction of the presenter device's create-infos, pinned for /// the lifetime of the `pyrowave_device` (heap boxes; moving `Hold` moves only pointers). struct Hold { @@ -174,6 +243,153 @@ struct PlaneSet { initialized: bool, } +/// A plane ring superseded by a mid-stream resize, awaiting safe destruction (see +/// [`RETIRE_HANDOVERS`] for the lifetime argument). +struct RetiredRing { + sets: Vec, + /// Frames handed to the presenter since this ring was retired. + handed_over: u32, + retired_at: Instant, +} + +/// One decode-output plane: R8, storage (decode writes) + sampled (presenter CSC). +unsafe fn make_plane( + device: &ash::Device, + mem_props: &vk::PhysicalDeviceMemoryProperties, + w: u32, + h: u32, +) -> Result<(vk::Image, vk::DeviceMemory, vk::ImageView)> { + let img = device.create_image( + &vk::ImageCreateInfo::default() + .image_type(vk::ImageType::TYPE_2D) + .format(vk::Format::R8_UNORM) + .extent(vk::Extent3D { + width: w, + height: h, + depth: 1, + }) + .mip_levels(1) + .array_layers(1) + .samples(vk::SampleCountFlags::TYPE_1) + .tiling(vk::ImageTiling::OPTIMAL) + .usage(vk::ImageUsageFlags::STORAGE | vk::ImageUsageFlags::SAMPLED) + .initial_layout(vk::ImageLayout::UNDEFINED), + None, + )?; + let req = device.get_image_memory_requirements(img); + let ti = (0..mem_props.memory_type_count) + .find(|&i| { + (req.memory_type_bits & (1 << i)) != 0 + && mem_props.memory_types[i as usize] + .property_flags + .contains(vk::MemoryPropertyFlags::DEVICE_LOCAL) + }) + .unwrap_or(0); + let mem = match device.allocate_memory( + &vk::MemoryAllocateInfo::default() + .allocation_size(req.size) + .memory_type_index(ti), + None, + ) { + Ok(m) => m, + Err(e) => { + device.destroy_image(img, None); + return Err(e.into()); + } + }; + if let Err(e) = device.bind_image_memory(img, mem, 0) { + device.destroy_image(img, None); + device.free_memory(mem, None); + return Err(e.into()); + } + let view = match device.create_image_view( + &vk::ImageViewCreateInfo::default() + .image(img) + .view_type(vk::ImageViewType::TYPE_2D) + .format(vk::Format::R8_UNORM) + .subresource_range(vk::ImageSubresourceRange { + aspect_mask: vk::ImageAspectFlags::COLOR, + base_mip_level: 0, + level_count: 1, + base_array_layer: 0, + layer_count: 1, + }), + None, + ) { + Ok(v) => v, + Err(e) => { + device.destroy_image(img, None); + device.free_memory(mem, None); + return Err(e.into()); + } + }; + Ok((img, mem, view)) +} + +unsafe fn destroy_sets(device: &ash::Device, sets: &[PlaneSet]) { + for set in sets { + for v in set.views { + device.destroy_image_view(v, None); + } + for i in set.imgs { + device.destroy_image(i, None); + } + for m in set.mems { + device.free_memory(m, None); + } + } +} + +/// Build a fresh [`RING`]-deep plane ring at the given dimensions; cleans up the partial +/// ring on failure (the caller keeps whatever it was using before). +unsafe fn build_ring( + device: &ash::Device, + mem_props: &vk::PhysicalDeviceMemoryProperties, + width: u32, + height: u32, +) -> Result> { + let mut ring: Vec = Vec::with_capacity(RING); + for _ in 0..RING { + let built = (|| -> Result { + let (y, ym, yv) = make_plane(device, mem_props, width, height)?; + let (cb, cbm, cbv) = match make_plane(device, mem_props, width / 2, height / 2) { + Ok(p) => p, + Err(e) => { + device.destroy_image_view(yv, None); + device.destroy_image(y, None); + device.free_memory(ym, None); + return Err(e); + } + }; + let (cr, crm, crv) = match make_plane(device, mem_props, width / 2, height / 2) { + Ok(p) => p, + Err(e) => { + for (v, i, m) in [(yv, y, ym), (cbv, cb, cbm)] { + device.destroy_image_view(v, None); + device.destroy_image(i, None); + device.free_memory(m, None); + } + return Err(e); + } + }; + Ok(PlaneSet { + imgs: [y, cb, cr], + mems: [ym, cbm, crm], + views: [yv, cbv, crv], + initialized: false, + }) + })(); + match built { + Ok(set) => ring.push(set), + Err(e) => { + destroy_sets(device, &ring); + return Err(e); + } + } + } + Ok(ring) +} + pub struct PyroWaveDecoder { // ash wrappers reconstructed over the presenter's raw handles (not owned — the // presenter outlives the decoder; Drop destroys only what this struct created). @@ -184,10 +400,13 @@ pub struct PyroWaveDecoder { pw_dev: pw::pyrowave_device, pw_dec: pw::pyrowave_decoder, ring: Vec, + /// Plane rings superseded by mid-stream resizes, pending safe destruction. + retired: Vec, next: usize, cmd_pool: vk::CommandPool, cmd: vk::CommandBuffer, fence: vk::Fence, + mem_props: vk::PhysicalDeviceMemoryProperties, width: u32, height: u32, /// The wire shard payload — the parse-window size for chunk-aligned AUs (§4.4): each @@ -294,68 +513,14 @@ impl PyroWaveDecoder { let mem_props = instance.get_physical_device_memory_properties( vk::PhysicalDevice::from_raw(vkd.physical_device as u64), ); - let make_plane = |w: u32, h: u32| -> Result<(vk::Image, vk::DeviceMemory, vk::ImageView)> { - let img = device.create_image( - &vk::ImageCreateInfo::default() - .image_type(vk::ImageType::TYPE_2D) - .format(vk::Format::R8_UNORM) - .extent(vk::Extent3D { - width: w, - height: h, - depth: 1, - }) - .mip_levels(1) - .array_layers(1) - .samples(vk::SampleCountFlags::TYPE_1) - .tiling(vk::ImageTiling::OPTIMAL) - .usage(vk::ImageUsageFlags::STORAGE | vk::ImageUsageFlags::SAMPLED) - .initial_layout(vk::ImageLayout::UNDEFINED), - None, - )?; - let req = device.get_image_memory_requirements(img); - let ti = (0..mem_props.memory_type_count) - .find(|&i| { - (req.memory_type_bits & (1 << i)) != 0 - && mem_props.memory_types[i as usize] - .property_flags - .contains(vk::MemoryPropertyFlags::DEVICE_LOCAL) - }) - .unwrap_or(0); - let mem = device.allocate_memory( - &vk::MemoryAllocateInfo::default() - .allocation_size(req.size) - .memory_type_index(ti), - None, - )?; - device.bind_image_memory(img, mem, 0)?; - let view = device.create_image_view( - &vk::ImageViewCreateInfo::default() - .image(img) - .view_type(vk::ImageViewType::TYPE_2D) - .format(vk::Format::R8_UNORM) - .subresource_range(vk::ImageSubresourceRange { - aspect_mask: vk::ImageAspectFlags::COLOR, - base_mip_level: 0, - level_count: 1, - base_array_layer: 0, - layer_count: 1, - }), - None, - )?; - Ok((img, mem, view)) + let ring = match build_ring(&device, &mem_props, width, height) { + Ok(r) => r, + Err(e) => { + pw::pyrowave_decoder_destroy(pw_dec); + pw::pyrowave_device_destroy(pw_dev); + return Err(e); + } }; - let mut ring = Vec::with_capacity(RING); - for _ in 0..RING { - let (y, ym, yv) = make_plane(width, height)?; - let (cb, cbm, cbv) = make_plane(width / 2, height / 2)?; - let (cr, crm, crv) = make_plane(width / 2, height / 2)?; - ring.push(PlaneSet { - imgs: [y, cb, cr], - mems: [ym, cbm, crm], - views: [yv, cbv, crv], - initialized: false, - }); - } let cmd_pool = device.create_command_pool( &vk::CommandPoolCreateInfo::default() @@ -383,16 +548,94 @@ impl PyroWaveDecoder { pw_dev, pw_dec, ring, + retired: Vec::new(), next: 0, cmd_pool, cmd, fence, + mem_props, width, height, wire_window: shard_payload.max(64), }) } + /// Mid-stream resize: rebuild the pyrowave decoder + plane ring at the new + /// dimensions in place, keeping the (dimension-independent) pyrowave device, command + /// pool, fence and pinned create-infos. Build-new-before-drop-old: a failure leaves + /// the current decoder untouched (and propagates — with the stream now at a size we + /// can't decode, the session ends with a real error instead of a frozen picture). + /// The old ring is RETIRED, not destroyed: the presenter / frame channel may still + /// reference its views (see [`RETIRE_HANDOVERS`]). + unsafe fn reconfigure(&mut self, width: u32, height: u32) -> Result<()> { + if width % 2 != 0 || height % 2 != 0 { + bail!("pyrowave 4:2:0 needs even dimensions (resize to {width}x{height})"); + } + let dinfo = pw::pyrowave_decoder_create_info { + device: self.pw_dev, + width: width as i32, + height: height as i32, + chroma: pw::pyrowave_chroma_subsampling_PYROWAVE_CHROMA_SUBSAMPLING_420, + fragment_path: false, + }; + let mut new_dec: pw::pyrowave_decoder = std::ptr::null_mut(); + pw_check( + pw::pyrowave_decoder_create(&dinfo, &mut new_dec), + "decoder_create (mid-stream resize)", + )?; + let new_ring = match build_ring(&self.device, &self.mem_props, width, height) { + Ok(r) => r, + Err(e) => { + pw::pyrowave_decoder_destroy(new_dec); + return Err(e).context("plane ring (mid-stream resize)"); + } + }; + // Our own decode work is fence-synchronous (never in flight here), so the old + // pyrowave decoder can go immediately; only the plane images wait (retired). + pw::pyrowave_decoder_destroy(self.pw_dec); + self.pw_dec = new_dec; + let old = std::mem::replace(&mut self.ring, new_ring); + self.retired.push(RetiredRing { + sets: old, + handed_over: 0, + retired_at: Instant::now(), + }); + self.next = 0; + tracing::info!( + from = %format!("{}x{}", self.width, self.height), + to = %format!("{width}x{height}"), + "PyroWave decoder rebuilt for mid-stream resize" + ); + self.width = width; + self.height = height; + Ok(()) + } + + /// Destroy retired rings that are provably unreachable (enough new-ring frames handed + /// over + a wall-clock floor — see [`RETIRE_HANDOVERS`]); the queue idle bounds any + /// still-submitted presenter sampling of the retiring views. + unsafe fn reap_retired(&mut self) { + let ripe = |r: &RetiredRing| { + r.handed_over >= RETIRE_HANDOVERS && r.retired_at.elapsed() >= RETIRE_MIN_AGE + }; + if !self.retired.iter().any(ripe) { + return; + } + { + let _guard = self.queue_lock.guard(); + let _ = self.device.queue_wait_idle(self.queue); + } + let mut kept = Vec::new(); + for r in self.retired.drain(..) { + if ripe(&r) { + destroy_sets(&self.device, &r.sets); + } else { + kept.push(r); + } + } + self.retired = kept; + } + /// One AU in → one frame out. `aligned` = the AU is shard-window chunked (each /// `wire_window` holds whole self-delimiting packets, zero-padded — walk and strip); /// `complete` = every shard arrived (a partial decodes anyway: missing blocks are @@ -477,20 +720,43 @@ impl PyroWaveDecoder { aligned: bool, complete: bool, ) -> Result> { + // Mid-stream resize: every frame's bitstream opens with a sequence header + // carrying its dimensions, so the AU itself announces the host's mode switch — + // no control-plane ordering to race (the Reconfigured ack travels on another + // stream). Upstream hard-errors on a dimension mismatch, so rebuild FIRST. A + // partial that lost its first shard sniffs `None` and decodes at the current + // size (correct when the size didn't change; harmlessly dropped below when it + // did — the next complete frame carries the header again). + if let Some(dims) = au_dims(au, aligned, self.wire_window) { + if dims != (self.width, self.height) { + self.reconfigure(dims.0, dims.1)?; + } + } + let mut push_err: Option = None; if aligned { let mut frag: Vec = Vec::new(); for win in au.chunks(self.wire_window) { - self.push_window(win, &mut frag)?; + if let Err(e) = self.push_window(win, &mut frag) { + push_err = Some(e); + break; + } } - } else { - pw_check( - pw::pyrowave_decoder_push_packet( - self.pw_dec, - au.as_ptr() as *const c_void, - au.len(), - ), - "push_packet", - )?; + } else if let Err(e) = pw_check( + pw::pyrowave_decoder_push_packet(self.pw_dec, au.as_ptr() as *const c_void, au.len()), + "push_packet", + ) { + push_err = Some(e); + } + if let Some(e) = push_err { + // A partial straddling a resize can carry blocks the (possibly wrong-size) + // decoder rejects — that's one lost frame, not a broken session; the next + // complete frame re-anchors (all-intra). A COMPLETE frame that fails to + // parse is real corruption: propagate. + if complete { + return Err(e); + } + tracing::debug!(error = %format!("{e:#}"), "partial AU rejected — frame dropped"); + return Ok(None); } // A complete AU that isn't ready is a stale/duplicate (sequence rewind) — skip. // A PARTIAL is decoded regardless: missing wavelet blocks reconstruct as zeros, @@ -604,6 +870,13 @@ impl PyroWaveDecoder { .context("pyrowave decode fence")?; self.ring[slot].initialized = true; + // This frame is about to reach the presenter — it advances every retired ring's + // displacement count, and ripe rings can now be destroyed. + for r in &mut self.retired { + r.handed_over += 1; + } + self.reap_retired(); + Ok(Some(PyroWavePlanarFrame { views: [ self.ring[slot].views[0].as_raw(), @@ -639,16 +912,9 @@ impl Drop for PyroWaveDecoder { } pw::pyrowave_decoder_destroy(self.pw_dec); pw::pyrowave_device_destroy(self.pw_dev); - for set in &self.ring { - for v in set.views { - self.device.destroy_image_view(v, None); - } - for i in set.imgs { - self.device.destroy_image(i, None); - } - for m in set.mems { - self.device.free_memory(m, None); - } + destroy_sets(&self.device, &self.ring); + for r in &self.retired { + destroy_sets(&self.device, &r.sets); } self.device.destroy_fence(self.fence, None); self.device.destroy_command_pool(self.cmd_pool, None); @@ -656,3 +922,104 @@ impl Drop for PyroWaveDecoder { } } } + +#[cfg(test)] +mod tests { + use super::{au_dims, seq_header_dims}; + + /// Little-endian encoding of upstream's `BitstreamSequenceHeader` bitfields (see + /// pyrowave_common.hpp): word 0 = width_minus_1:14 | height_minus_1:14 | sequence:3 + /// | extended:1; word 1 = total_blocks:24 | code:2 | chroma:1 | … + fn seq_header(w: u32, h: u32, code: u32) -> [u8; 8] { + let w0 = (w - 1) & 0x3FFF | ((h - 1) & 0x3FFF) << 14 | 1 << 31; + let w1 = 0x1234 | code << 24; // arbitrary total_blocks + let mut out = [0u8; 8]; + out[0..4].copy_from_slice(&w0.to_le_bytes()); + out[4..8].copy_from_slice(&w1.to_le_bytes()); + out + } + + /// A regular `BitstreamHeader` (block packet): extended bit clear. + fn block_header() -> [u8; 8] { + let w0 = 0xBEEFu32 | 8 << 16; // ballot | payload_words=8, extended=0 + let w1 = 42u32 << 8; // block_index + let mut out = [0u8; 8]; + out[0..4].copy_from_slice(&w0.to_le_bytes()); + out[4..8].copy_from_slice(&w1.to_le_bytes()); + out + } + + /// Wrap `body` in one §4.4 framed window of `win` bytes (4-byte prefix + zero pad). + fn window(body: &[u8], kind: u16, win: usize) -> Vec { + let mut out = Vec::with_capacity(win); + out.extend_from_slice(&(body.len() as u16).to_le_bytes()); + out.extend_from_slice(&kind.to_le_bytes()); + out.extend_from_slice(body); + out.resize(win, 0); + out + } + + #[test] + fn sniffs_dims_from_a_sequence_header() { + assert_eq!( + seq_header_dims(&seq_header(1920, 1080, 0)), + Some((1920, 1080)) + ); + assert_eq!( + seq_header_dims(&seq_header(1280, 720, 0)), + Some((1280, 720)) + ); + // 14-bit fields carry up to 16384. + assert_eq!( + seq_header_dims(&seq_header(16384, 16384, 0)), + Some((16384, 16384)) + ); + } + + #[test] + fn rejects_non_sequence_headers() { + assert_eq!(seq_header_dims(&block_header()), None); // extended bit clear + assert_eq!(seq_header_dims(&seq_header(1920, 1080, 1)), None); // not START_OF_FRAME + assert_eq!(seq_header_dims(&seq_header(1920, 1080, 0)[..7]), None); // short + assert_eq!(seq_header_dims(&[]), None); + } + + #[test] + fn unaligned_au_sniffs_at_byte_zero() { + let mut au = seq_header(2560, 1440, 0).to_vec(); + au.extend_from_slice(&block_header()); + assert_eq!(au_dims(&au, false, 1404), Some((2560, 1440))); + } + + #[test] + fn aligned_au_sniffs_the_first_window_body() { + const WIN: usize = 64; + let mut body = seq_header(1280, 800, 0).to_vec(); + body.extend_from_slice(&block_header()); + // WIN_PACKED first window, then another window of blocks. + let mut au = window(&body, 0, WIN); + au.extend_from_slice(&window(&block_header(), 0, WIN)); + assert_eq!(au_dims(&au, true, WIN), Some((1280, 800))); + // An oversized first packet rides a FRAG chain — FRAG_FIRST also starts at the + // frame's first byte, so the header is still there. + let frag = window(&body, 1, WIN); + assert_eq!(au_dims(&frag, true, WIN), Some((1280, 800))); + } + + #[test] + fn lost_first_window_means_unknown_dims() { + const WIN: usize = 64; + // A lost shard arrives as a zeroed window (used = 0) — the sequence header is gone. + let mut au = vec![0u8; WIN]; + au.extend_from_slice(&window(&seq_header(1280, 800, 0), 0, WIN)); + assert_eq!(au_dims(&au, true, WIN), None); + // A FRAG_CONT/LAST first window means the same (its FIRST was in a lost prior AU). + let cont = window(&block_header(), 2, WIN); + assert_eq!(au_dims(&cont, true, WIN), None); + // Garbage used-length never reads out of bounds. + let mut garbage = vec![0u8; WIN]; + garbage[0] = 0xFF; + garbage[1] = 0xFF; + assert_eq!(au_dims(&garbage, true, WIN), None); + } +} diff --git a/crates/punktfunk-host/src/encode.rs b/crates/punktfunk-host/src/encode.rs index 7049e831..e3ba51bf 100644 --- a/crates/punktfunk-host/src/encode.rs +++ b/crates/punktfunk-host/src/encode.rs @@ -428,6 +428,15 @@ pub fn validate_dimensions(codec: Codec, width: u32, height: u32) -> Result<()> if width % 2 != 0 || height % 2 != 0 { anyhow::bail!("invalid encode resolution {width}x{height}: dimensions must be even"); } + // PyroWave's 5-level wavelet decomposition needs ≥ 4·2⁵ px per axis (upstream + // `MinimumImageSize` — the band mirroring breaks below it); reject a tiny mode here + // (e.g. a match-window resize dragged to a sliver) instead of failing the encoder + // rebuild after the switch was acked. + if codec == Codec::PyroWave && (width < 128 || height < 128) { + anyhow::bail!( + "invalid PyroWave resolution {width}x{height}: the wavelet needs at least 128px per axis" + ); + } let max = codec.max_dimension(); if width > max || height > max { anyhow::bail!( diff --git a/crates/punktfunk-host/src/encode/linux/pyrowave.rs b/crates/punktfunk-host/src/encode/linux/pyrowave.rs index cf68f6f8..67749b12 100644 --- a/crates/punktfunk-host/src/encode/linux/pyrowave.rs +++ b/crates/punktfunk-host/src/encode/linux/pyrowave.rs @@ -1172,7 +1172,7 @@ mod tests { buf.row_stride_in_bytes = [w as usize, (w / 2) as usize, (w / 2) as usize]; buf.plane_size_in_bytes = [y.len(), cb.len(), cr.len()]; assert_eq!( - pw::pyrowave_decoder_decode_cpu_buffer_synchronous(dec, &mut buf), + pw::pyrowave_decoder_decode_cpu_buffer_synchronous(dec, &buf), pw::pyrowave_result_PYROWAVE_SUCCESS ); pw::pyrowave_decoder_destroy(dec); diff --git a/crates/punktfunk-host/src/punktfunk1.rs b/crates/punktfunk-host/src/punktfunk1.rs index a9ff0690..7871d333 100644 --- a/crates/punktfunk-host/src/punktfunk1.rs +++ b/crates/punktfunk-host/src/punktfunk1.rs @@ -1606,6 +1606,9 @@ async fn serve_session( } }); let bitrate_kbps = welcome.bitrate_kbps; // resolved encoder bitrate (Hello clamped, or default) + // "Automatic" request: the resolved rate is a host default — for PyroWave a per-mode + // bpp pin the data plane re-resolves on a mid-stream mode switch. + let bitrate_auto = hello.bitrate_kbps == 0; let bit_depth = welcome.bit_depth; // resolved encode bit depth (8, or 10 when negotiated) // Resolved chroma — derive the typed value back from the wire byte the Welcome carried (so the // session uses exactly what the client was told). `Yuv444` only when the handshake gate passed. @@ -1703,6 +1706,7 @@ async fn serve_session( bitrate_rx, compositor, bitrate_kbps, + bitrate_auto, bit_depth, chroma, codec, @@ -3948,6 +3952,11 @@ struct SessionContext { compositor: crate::vdisplay::Compositor, /// Negotiated encoder bitrate (kbps). bitrate_kbps: u32, + /// The client asked for "Automatic" (`Hello::bitrate_kbps == 0`), so `bitrate_kbps` came from + /// the host's codec-aware default. For PyroWave that default is the ~1.6 bpp operating point of + /// the NEGOTIATED MODE (`resolve_bitrate_kbps_for`) — a mid-stream mode switch re-resolves it + /// for the new mode (the pin follows the resolution; an explicit client rate stays put). + bitrate_auto: bool, /// Negotiated encode bit depth (8, or 10 = HEVC Main10). bit_depth: u8, /// Negotiated chroma subsampling (4:2:0, or 4:4:4 when the client + host + GPU all support it). @@ -4027,6 +4036,7 @@ fn virtual_stream(ctx: SessionContext) -> Result<()> { bitrate_rx, compositor, mut bitrate_kbps, + bitrate_auto, bit_depth, // The resolved chroma is already captured in `plan` (above); ignore the duplicate here. chroma: _, @@ -4363,11 +4373,29 @@ fn virtual_stream(ctx: SessionContext) -> Result<()> { } if let Some(new_mode) = want { tracing::info!(?new_mode, "rebuilding pipeline for mode switch"); + // PyroWave's Automatic bitrate is a per-mode ~1.6 bpp pin (resolve_bitrate_kbps_for) — + // a resolution change moves the operating point (1080p→4K quadruples the pixel rate), + // so re-resolve it for the new mode. Explicit client rates stay put (the operator knows + // the link), and the H.26x codecs keep their mode-independent rate (ABR owns it). + let mode_bitrate = if bitrate_auto && plan.codec == crate::encode::Codec::PyroWave { + resolve_bitrate_kbps_for(plan.codec, 0, &new_mode) + } else { + bitrate_kbps + }; // Build the new pipeline BEFORE dropping the old one: the host already acked // the switch as accepted, so a rebuild failure must not kill an otherwise // healthy session — keep streaming the current mode and log instead. - match build_pipeline(&mut vd, new_mode, bitrate_kbps, bit_depth, plan, &quit) { + match build_pipeline(&mut vd, new_mode, mode_bitrate, bit_depth, plan, &quit) { Ok(next_pipe) => { + if mode_bitrate != bitrate_kbps { + tracing::info!( + from_kbps = bitrate_kbps, + to_kbps = mode_bitrate, + "pinned PyroWave bitrate re-resolved for the new mode" + ); + bitrate_kbps = mode_bitrate; + live_bitrate.store(mode_bitrate, Ordering::Relaxed); + } let old_display_gen = cur_display_gen; // The destructuring assignment drops the OLD capturer (→ its display lease) as // each binding is replaced — the new pipeline is already up (create-before-drop).