feat(client,host): PyroWave Apple Metal decoder + per-mode bitrate pin

- clients/apple: native Metal wavelet decoder + compute shaders (Phase 5),
  decoding PyroWave without embedding MoltenVK.
- pf-client-core: plumb user_flags/completeness through Decoder::decode_frame
  so the PyroWave backend parses chunk-aligned + partial AUs; gate the param's
  unused-warning to exactly the non-pyrowave builds (fixes -D warnings on the
  featureless Linux client build).
- punktfunk-host: on a mid-stream mode switch, re-resolve the "Automatic"
  PyroWave bitrate for the new mode's ~1.6 bpp operating point (explicit rates
  and H.26x ABR stay put); reject sub-128px PyroWave modes before the encoder
  rebuild instead of after the ack.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-15 11:47:42 +02:00
parent 2621b6e6b1
commit 9127c3465f
8 changed files with 1713 additions and 83 deletions
@@ -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
// writesampled-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 40,
/// component 02 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 YCbCrRGB 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<UInt8>) -> 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..<payloadWords {
payload.append(loadWord(base, pos + w * 4))
}
}
} else if layout != nil {
return false // out-of-bounds block index corrupt stream
}
// No layout yet (SOF lost): skip the packet, the block stays missing.
pos += payloadWords * 4
}
// In the windowed framing, `used` delimits exactly; dense AUs must also consume
// fully (upstream errors on trailing bytes).
return pos == count
}
private func loadWord(_ base: UnsafePointer<UInt8>, _ 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 01) / R32Float (levels 24)
/// 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..<WaveletLayout.decompositionLevels {
for component in 0..<3 {
if level == 0 && component != 0 { continue } // 4:2:0
for band in (level == WaveletLayout.decompositionLevels - 1 ? 0 : 1)..<4 {
let meta = layout.blockMeta[component][level][band]
let w = layout.levelWidth(level)
let h = layout.levelHeight(level)
var regs = DequantRegisters(
resolution: SIMD2(Int32(w), Int32(h)),
outputLayer: Int32(band),
blockOffset32x32: Int32(meta.offset),
blockStride32x32: Int32(meta.stride))
dequant.setTexture(coefficients[component][level], index: 0)
dequant.setBytes(
&regs, length: MemoryLayout<DequantRegisters>.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 writesampled-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(&regs, length: MemoryLayout<IdwtRegisters>.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..<WaveletLayout.decompositionLevels {
let desc = MTLTextureDescriptor()
desc.textureType = .type2DArray
desc.arrayLength = 4
// Upstream precision 1: fp16 storage for the two finest levels, fp32 for the
// coarse levels whose values feed every later reconstruction step.
desc.pixelFormat = level < 2 ? .r16Float : .r32Float
desc.width = newLayout.levelWidth(level)
desc.height = newLayout.levelHeight(level)
desc.usage = [.shaderRead, .shaderWrite]
desc.storageMode = .private
guard let tex = device.makeTexture(descriptor: desc) else { return false }
tex.label = "pyrowave coeff c\(component) L\(level)"
guard
let ll = tex.makeTextureView(
pixelFormat: desc.pixelFormat, textureType: .type2D,
levels: 0..<1, slices: 0..<1)
else { return false }
ll.label = "pyrowave LL c\(component) L\(level)"
perLevel.append(tex)
perLevelLL.append(ll)
}
coeff.append(perLevel)
lls.append(perLevelLL)
}
var newSlots: [Slot] = []
for i in 0..<Self.ringDepth {
let plane = { (w: Int, h: Int, name: String) -> 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<Int32>
var outputLayer: Int32
var blockOffset32x32: Int32
var blockStride32x32: Int32
}
private struct IdwtRegisters {
var resolution: SIMD2<Int32>
var invResolution: SIMD2<Float>
}
}
#endif
@@ -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 01 and
// R32Float for levels 24 (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 <metal_stdlib>
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<float, access::write> uDequantImg [[texture(0)]],
const device uint *payload_offsets [[buffer(0)]],
const device uint *payload_u32 [[buffer(1)]],
constant DequantRegisters &registers [[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<const device ushort *>(payload_u32);
const device uchar *payload_u8 = reinterpret_cast<const device uchar *>(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<float, access::sample> 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<float, access::sample> 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<float, access::sample> uTexture [[texture(0)]],
texture2d<float, access::write> uOutput [[texture(1)]],
sampler uSampler [[sampler(0)]],
constant IdwtRegisters &registers [[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));
}
}
}
"""
+66
View File
@@ -328,6 +328,20 @@ fn pump(
// Live host↔client clock offset: loaded per frame (Relaxed) so mid-stream re-syncs (an NTP // 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. // step, drift) keep the capture-clock latency stats honest — never cached at session start.
let clock_offset_live = connector.clock_offset_shared(); 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; let mut total_frames = 0u64;
// Newest frame index handed to the decoder — the staleness bar for late partials. // Newest frame index handed to the decoder — the staleness bar for late partials.
let mut newest_decoded_idx: Option<u32> = None; let mut newest_decoded_idx: Option<u32> = None;
@@ -368,6 +382,18 @@ fn pump(
if stop.load(Ordering::SeqCst) { if stop.load(Ordering::SeqCst) {
break None; 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 // 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 // responsiveness and the per-iteration keyframe-recovery check (a frame arrives
// every ~816 ms at 60120 Hz anyway, so this rarely times out mid-stream). // every ~816 ms at 60120 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")) .map_err(|e| tracing::warn!(error = %e, "audio thread failed to start — audio disabled"))
.ok() .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");
}
}
}
+5
View File
@@ -628,6 +628,11 @@ impl Decoder {
pub fn decode_frame( pub fn decode_frame(
&mut self, &mut self,
au: &[u8], 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, user_flags: u32,
complete: bool, complete: bool,
) -> Result<Option<DecodedImage>> { ) -> Result<Option<DecodedImage>> {
+447 -80
View File
@@ -19,6 +19,14 @@
//! content-equivalent ones from [`VulkanDecodeDevice`]'s exported extension lists, //! content-equivalent ones from [`VulkanDecodeDevice`]'s exported extension lists,
//! feature facts and queue-family shape (pyrowave reads them for extension/feature //! feature facts and queue-family shape (pyrowave reads them for extension/feature
//! detection; pointer identity is not required). //! 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 crate::video::{ColorDesc, VulkanDecodeDevice};
use anyhow::{bail, Context as _, Result}; use anyhow::{bail, Context as _, Result};
@@ -27,6 +35,7 @@ use ash::vk::Handle as _;
use pyrowave_sys as pw; use pyrowave_sys as pw;
use std::ffi::{c_char, c_void, CString}; use std::ffi::{c_char, c_void, CString};
use std::sync::Arc; use std::sync::Arc;
use std::time::{Duration, Instant};
/// Plane-set ring depth: decode writes slot N while the presenter may still sample /// 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 /// 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). /// keeps LOGICAL reuse far enough behind).
const RING: usize = 4; 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<()> { fn pw_check(r: pw::pyrowave_result, what: &str) -> Result<()> {
if r == pw::pyrowave_result_PYROWAVE_SUCCESS { if r == pw::pyrowave_result_PYROWAVE_SUCCESS {
Ok(()) 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 /// 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). /// the lifetime of the `pyrowave_device` (heap boxes; moving `Hold` moves only pointers).
struct Hold { struct Hold {
@@ -174,6 +243,153 @@ struct PlaneSet {
initialized: bool, 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<PlaneSet>,
/// 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<Vec<PlaneSet>> {
let mut ring: Vec<PlaneSet> = Vec::with_capacity(RING);
for _ in 0..RING {
let built = (|| -> Result<PlaneSet> {
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 { pub struct PyroWaveDecoder {
// ash wrappers reconstructed over the presenter's raw handles (not owned — the // ash wrappers reconstructed over the presenter's raw handles (not owned — the
// presenter outlives the decoder; Drop destroys only what this struct created). // presenter outlives the decoder; Drop destroys only what this struct created).
@@ -184,10 +400,13 @@ pub struct PyroWaveDecoder {
pw_dev: pw::pyrowave_device, pw_dev: pw::pyrowave_device,
pw_dec: pw::pyrowave_decoder, pw_dec: pw::pyrowave_decoder,
ring: Vec<PlaneSet>, ring: Vec<PlaneSet>,
/// Plane rings superseded by mid-stream resizes, pending safe destruction.
retired: Vec<RetiredRing>,
next: usize, next: usize,
cmd_pool: vk::CommandPool, cmd_pool: vk::CommandPool,
cmd: vk::CommandBuffer, cmd: vk::CommandBuffer,
fence: vk::Fence, fence: vk::Fence,
mem_props: vk::PhysicalDeviceMemoryProperties,
width: u32, width: u32,
height: u32, height: u32,
/// The wire shard payload — the parse-window size for chunk-aligned AUs (§4.4): each /// 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( let mem_props = instance.get_physical_device_memory_properties(
vk::PhysicalDevice::from_raw(vkd.physical_device as u64), vk::PhysicalDevice::from_raw(vkd.physical_device as u64),
); );
let make_plane = |w: u32, h: u32| -> Result<(vk::Image, vk::DeviceMemory, vk::ImageView)> { let ring = match build_ring(&device, &mem_props, width, height) {
let img = device.create_image( Ok(r) => r,
&vk::ImageCreateInfo::default() Err(e) => {
.image_type(vk::ImageType::TYPE_2D) pw::pyrowave_decoder_destroy(pw_dec);
.format(vk::Format::R8_UNORM) pw::pyrowave_device_destroy(pw_dev);
.extent(vk::Extent3D { return Err(e);
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 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( let cmd_pool = device.create_command_pool(
&vk::CommandPoolCreateInfo::default() &vk::CommandPoolCreateInfo::default()
@@ -383,16 +548,94 @@ impl PyroWaveDecoder {
pw_dev, pw_dev,
pw_dec, pw_dec,
ring, ring,
retired: Vec::new(),
next: 0, next: 0,
cmd_pool, cmd_pool,
cmd, cmd,
fence, fence,
mem_props,
width, width,
height, height,
wire_window: shard_payload.max(64), 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 /// 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); /// `wire_window` holds whole self-delimiting packets, zero-padded — walk and strip);
/// `complete` = every shard arrived (a partial decodes anyway: missing blocks are /// `complete` = every shard arrived (a partial decodes anyway: missing blocks are
@@ -477,20 +720,43 @@ impl PyroWaveDecoder {
aligned: bool, aligned: bool,
complete: bool, complete: bool,
) -> Result<Option<PyroWavePlanarFrame>> { ) -> Result<Option<PyroWavePlanarFrame>> {
// 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<anyhow::Error> = None;
if aligned { if aligned {
let mut frag: Vec<u8> = Vec::new(); let mut frag: Vec<u8> = Vec::new();
for win in au.chunks(self.wire_window) { 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( } else if let Err(e) = pw_check(
pw::pyrowave_decoder_push_packet( pw::pyrowave_decoder_push_packet(self.pw_dec, au.as_ptr() as *const c_void, au.len()),
self.pw_dec,
au.as_ptr() as *const c_void,
au.len(),
),
"push_packet", "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 complete AU that isn't ready is a stale/duplicate (sequence rewind) — skip.
// A PARTIAL is decoded regardless: missing wavelet blocks reconstruct as zeros, // A PARTIAL is decoded regardless: missing wavelet blocks reconstruct as zeros,
@@ -604,6 +870,13 @@ impl PyroWaveDecoder {
.context("pyrowave decode fence")?; .context("pyrowave decode fence")?;
self.ring[slot].initialized = true; 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 { Ok(Some(PyroWavePlanarFrame {
views: [ views: [
self.ring[slot].views[0].as_raw(), self.ring[slot].views[0].as_raw(),
@@ -639,16 +912,9 @@ impl Drop for PyroWaveDecoder {
} }
pw::pyrowave_decoder_destroy(self.pw_dec); pw::pyrowave_decoder_destroy(self.pw_dec);
pw::pyrowave_device_destroy(self.pw_dev); pw::pyrowave_device_destroy(self.pw_dev);
for set in &self.ring { destroy_sets(&self.device, &self.ring);
for v in set.views { for r in &self.retired {
self.device.destroy_image_view(v, None); destroy_sets(&self.device, &r.sets);
}
for i in set.imgs {
self.device.destroy_image(i, None);
}
for m in set.mems {
self.device.free_memory(m, None);
}
} }
self.device.destroy_fence(self.fence, None); self.device.destroy_fence(self.fence, None);
self.device.destroy_command_pool(self.cmd_pool, 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<u8> {
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);
}
}
+9
View File
@@ -428,6 +428,15 @@ pub fn validate_dimensions(codec: Codec, width: u32, height: u32) -> Result<()>
if width % 2 != 0 || height % 2 != 0 { if width % 2 != 0 || height % 2 != 0 {
anyhow::bail!("invalid encode resolution {width}x{height}: dimensions must be even"); 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(); let max = codec.max_dimension();
if width > max || height > max { if width > max || height > max {
anyhow::bail!( anyhow::bail!(
@@ -1172,7 +1172,7 @@ mod tests {
buf.row_stride_in_bytes = [w as usize, (w / 2) as usize, (w / 2) as usize]; 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()]; buf.plane_size_in_bytes = [y.len(), cb.len(), cr.len()];
assert_eq!( 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_result_PYROWAVE_SUCCESS
); );
pw::pyrowave_decoder_destroy(dec); pw::pyrowave_decoder_destroy(dec);
+29 -1
View File
@@ -1606,6 +1606,9 @@ async fn serve_session(
} }
}); });
let bitrate_kbps = welcome.bitrate_kbps; // resolved encoder bitrate (Hello clamped, or default) 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) 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 // 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. // 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, bitrate_rx,
compositor, compositor,
bitrate_kbps, bitrate_kbps,
bitrate_auto,
bit_depth, bit_depth,
chroma, chroma,
codec, codec,
@@ -3948,6 +3952,11 @@ struct SessionContext {
compositor: crate::vdisplay::Compositor, compositor: crate::vdisplay::Compositor,
/// Negotiated encoder bitrate (kbps). /// Negotiated encoder bitrate (kbps).
bitrate_kbps: u32, 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). /// Negotiated encode bit depth (8, or 10 = HEVC Main10).
bit_depth: u8, bit_depth: u8,
/// Negotiated chroma subsampling (4:2:0, or 4:4:4 when the client + host + GPU all support it). /// 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, bitrate_rx,
compositor, compositor,
mut bitrate_kbps, mut bitrate_kbps,
bitrate_auto,
bit_depth, bit_depth,
// The resolved chroma is already captured in `plan` (above); ignore the duplicate here. // The resolved chroma is already captured in `plan` (above); ignore the duplicate here.
chroma: _, chroma: _,
@@ -4363,11 +4373,29 @@ fn virtual_stream(ctx: SessionContext) -> Result<()> {
} }
if let Some(new_mode) = want { if let Some(new_mode) = want {
tracing::info!(?new_mode, "rebuilding pipeline for mode switch"); 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 // 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 // the switch as accepted, so a rebuild failure must not kill an otherwise
// healthy session — keep streaming the current mode and log instead. // 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) => { 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; let old_display_gen = cur_display_gen;
// The destructuring assignment drops the OLD capturer (→ its display lease) as // The destructuring assignment drops the OLD capturer (→ its display lease) as
// each binding is replaced — the new pipeline is already up (create-before-drop). // each binding is replaced — the new pipeline is already up (create-before-drop).