unom/punktfunk
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 25 Wiki Activity

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

Browse Source 
- 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:
Enrico Bühler
2026-07-15 11:47:42 +02:00
parent 2621b6e6b1
commit 9127c3465f
8 changed files with 1713 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clients/apple/Sources/PunktfunkKit/Video/MetalWaveletShaders.swift
+551
View File
@@ -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 <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));
}
}
}
"""