diff --git a/libavcodec/ffv1enc_vulkan.c b/libavcodec/ffv1enc_vulkan.c index 14f7d2e5dc..2a05632210 100644 --- a/libavcodec/ffv1enc_vulkan.c +++ b/libavcodec/ffv1enc_vulkan.c @@ -1191,8 +1191,15 @@ static int init_encode_shader(AVCodecContext *avctx, VkSpecializationInfo *sl) FFVulkanShader *shd = &fv->enc; uint32_t wg_x = fv->ctx.ac != AC_GOLOMB_RICE ? CONTEXT_SIZE : 1; + uint32_t sg_size = 0; + /* The range coder encoder needs its workgroup to be a single subgroup */ + if (fv->ctx.ac != AC_GOLOMB_RICE && + (fv->s.subgroup_props.requiredSubgroupSizeStages & VK_SHADER_STAGE_COMPUTE_BIT) && + fv->s.subgroup_props.minSubgroupSize <= CONTEXT_SIZE && + fv->s.subgroup_props.maxSubgroupSize >= CONTEXT_SIZE) + sg_size = CONTEXT_SIZE; ff_vk_shader_load(shd, VK_SHADER_STAGE_COMPUTE_BIT, sl, - (uint32_t []) { wg_x, 1, 1 }, 0); + (uint32_t []) { wg_x, 1, 1 }, sg_size); ff_vk_shader_add_push_const(shd, 0, sizeof(FFv1ShaderParams), VK_SHADER_STAGE_COMPUTE_BIT); @@ -1476,7 +1483,7 @@ static av_cold int vulkan_encode_ffv1_init(AVCodecContext *avctx) !fv->force_pcm && fv->optimize_rct && !f->bayer; /* Init shader specialization consts */ - SPEC_LIST_CREATE(sl, 19, 19*sizeof(uint32_t)) + SPEC_LIST_CREATE(sl, 20, 20*sizeof(uint32_t)) SPEC_LIST_ADD(sl, 0, 32, RGB_LINECACHE); SPEC_LIST_ADD(sl, 1, 32, f->ec); ff_ffv1_vk_set_common_sl(avctx, f, sl, fv->s.frames->sw_format); @@ -1514,6 +1521,18 @@ static av_cold int vulkan_encode_ffv1_init(AVCodecContext *avctx) RET(init_indirect(avctx, fv->ctx.use32bit ? AV_PIX_FMT_RGBA128 : AV_PIX_FMT_RGBA64)); + /* Keep the probability states on-chip if they fit. Line-interleaved + * planes need all of their states at once. */ + if (f->ac != AC_GOLOMB_RICE) { + uint32_t lds_bytes = CONTEXT_SIZE*f->context_count[f->context_model]; + uint32_t lds_max = fv->s.props.properties.limits.maxComputeSharedMemorySize; + if (fv->is_rgb || f->bayer) + lds_bytes *= f->plane_count; + if (lds_bytes + 2048 > lds_max) + lds_bytes = 0; + SPEC_LIST_ADD(sl, 19, 32, lds_bytes); + } + /* Encode shader */ err = init_encode_shader(avctx, sl); if (err < 0) diff --git a/libavcodec/vulkan/ffv1_dec.comp.glsl b/libavcodec/vulkan/ffv1_dec.comp.glsl index 8e52ac2aeb..2c6d830cec 100644 --- a/libavcodec/vulkan/ffv1_dec.comp.glsl +++ b/libavcodec/vulkan/ffv1_dec.comp.glsl @@ -43,63 +43,260 @@ layout(set = 1, binding = 6) readonly buffer fltmap_buf { #ifndef GOLOMB +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_ballot : require + layout (set = 1, binding = 3, scalar) buffer slice_state_buf { uint8_t slice_rc_state[]; }; +layout (set = 1, binding = 3, scalar) buffer slice_state_buf32 { + uint32_t slice_rc_state32[]; +}; + +/* Size of the on-chip copy of the probability states, in bytes. Slices + * whose states do not fit access them from slice_rc_state instead. */ +layout (constant_id = 19) const uint lds_state_bytes = 0; + +shared uint lds_state[lds_state_bytes > 0 ? (lds_state_bytes >> 2) : 1]; +/* Both possible next states of each state, as zero | (one << 8) */ +shared uint16_t lds_zo[256]; + +layout (buffer_reference, buffer_reference_align = 4) readonly buffer u32buf_ro { + uint32_t v; +}; + +bool use_lds; + +/* + * Register-resident range decoder. + * The whole workgroup is a single subgroup, and every invocation performs + * the exact same computation on the same data, so all of this is + * dynamically uniform, and lives in scalar registers on hardware that has + * them. Nothing on the critical path of a decision touches memory. + */ +uint rl_low; +uint rl_range; +uint rl_off; /* Offset of the next byte to shift in */ +uint rl_end; +uint rl_cur; /* Bytes of the current word not yet shifted in */ +uint rl_cnt; /* Number of bytes left in rl_cur */ +uint rl_wpos; /* Offset of the next word */ + +/* 256 bytes of upcoming bitstream, spread across the subgroup, and loaded + * together with the rest of the data needed by the next 32 samples. */ +uint rl_wbase; +uint rl_w0; +uint rl_w1; + +uint rl_load(uint pos) +{ + return pos < rl_end ? u32buf_ro(slice_data + pos).v : 0; +} + +void rl_window(void) +{ + rl_wbase = rl_wpos; + uint pos = rl_wbase + 4*gl_SubgroupInvocationID; + rl_w0 = rl_load(pos); + rl_w1 = rl_load(pos + 128); +} + +void rl_next_word(void) +{ + uint o = rl_wpos - rl_wbase; + uint w; + if (expectEXT(o < 256, true)) + w = subgroupBroadcast(o < 128 ? rl_w0 : rl_w1, (o >> 2) & 31u); + else + w = rl_load(rl_wpos); + + /* Bytes past the end of the slice read as zero */ + if (expectEXT((rl_wpos + 4) > rl_end, false)) + w = rl_wpos >= rl_end ? 0 : (w & ((1u << ((rl_end - rl_wpos) << 3)) - 1u)); + + rl_cur = w; + rl_cnt = 4; + rl_wpos += 4; +} + +void rl_init(in RangeCoder c) +{ + rl_low = c.low; + rl_range = c.range; + rl_off = c.bs_off; + rl_end = c.bs_end; + + uint mis = uint(uint64_t(slice_data) + c.bs_off) & 3u; + rl_wpos = c.bs_off - mis; + rl_window(); + rl_next_word(); + rl_cur >>= mis << 3; + rl_cnt -= mis; +} + +void rl_refill(void) +{ + rl_range <<= 8; + rl_low = (rl_low << 8) | (rl_cur & 0xFFu); + rl_cur >>= 8; + rl_off++; + if (--rl_cnt == 0) + rl_next_word(); +} + +bool rl_get_internal(uint range1) +{ + uint ranged = rl_range - range1; + bool bit = rl_low >= ranged; + rl_low = bit ? rl_low - ranged : rl_low; + rl_range = bit ? range1 : ranged; + if (expectEXT(rl_range < 0x100, false)) + rl_refill(); + return bit; +} -#define READ(idx) get_rac_state(idx) -shared int sym_e; -shared bool rc_dec[CONTEXT_SIZE]; -int get_isymbol(void) +bool rl_get(uint state) { - sym_e = 0; - rc_dec[0] = true; - if (READ(0)) + return rl_get_internal((rl_range * state) >> 8); +} + +bool rl_get_equi(void) +{ + return rl_get_internal(rl_range >> 1); +} + +uint zo_next(uint state, bool bit) +{ + uint z = uint(lds_zo[state]); + return bit ? (z >> 8) : (z & 0xFFu); +} + +/* A 256-entry quantization table, spread across the subgroup: invocation i + * holds entries i + 32*k, k = 0..7, two per register. */ +uvec4 qt_load(uint qi, uint j) +{ + uint l = gl_SubgroupInvocationID; + uvec4 r; + [[unroll]] + for (int d = 0; d < 4; d++) + r[d] = uint(uint16_t(quant_table[qi][j][l + 64*d])) | + (uint(uint16_t(quant_table[qi][j][l + 64*d + 32])) << 16); + return r; +} + +int qt_get(uvec4 t, uint idx) +{ + uint d = idx >> 6; + uint v = subgroupBroadcast(d == 0 ? t.x : d == 1 ? t.y : d == 2 ? t.z : t.w, + idx & 31u); + return bitfieldExtract(int(v), int(idx & 32u) >> 1, 16); +} + +/* The 32 probability states of the current context: dword k holds states + * 4k..4k+3. The masks record which states were read, and with which + * result, so that they can all be adapted in parallel once the symbol has + * been decoded. */ +#define SBYTE(idx) bitfieldExtract((idx) < 16 ? sa[((idx) >> 2) & 3] : sb[((idx) >> 2) & 3], \ + ((idx) & 3)*8, 8) +uint sym_used; +uint sym_bits; + +int rl_get_isymbol(uvec4 sa, uvec4 sb, inout uint dw, inout uvec4 zt) +{ + const uint lane = gl_SubgroupInvocationID; + + bool b = rl_get(SBYTE(0)); + sym_used = 1u; + sym_bits = uint(b); + if (b) return 0; - int e = 1; - for (; e < 11; e++) { - rc_dec[e] = true; - if (!READ(e)) + /* Exponent, unary coded */ + uint e = 0; + uint sgs = SBYTE(21); + [[unroll]] + for (int k = 0; k < 10; k++) { + b = rl_get(SBYTE(1 + k)); + sym_used |= 1u << (1 + k); + sym_bits |= uint(b) << (1 + k); + if (!b) { + sgs = SBYTE(11 + k); break; + } + e++; } - int a = 1; - sym_e = e + 10; - rc_dec[sym_e] = true; - - if (c_bits > 10 && e == 11) { + if (expectEXT(b, false)) { + /* State 10 gets reused until a zero is read, and has to be adapted + * in between. The last read gets adapted along with the rest. */ + uint s10 = SBYTE(10); do { - rc_state[10] = zero_one_state[rc_state[10] + 256]; - e++; - } while (READ(10)); - - a = READ(31) ? 0x3 : 0x2; - for (e -= 2; e >= 11; e--) { - rc_state[31] = zero_one_state[rc_state[31] + - (rc_data[31] ? 256 : 0)]; - a <<= 1; - a |= int(READ(31)); + s10 = zo_next(s10, true); + b = rl_get(s10); + e += uint(b); + } while (b && e < 31); + sym_bits &= ~(1u << 10); + if (lane == 2) { + dw = bitfieldInsert(dw, s10, 16, 8); + zt.z = uint(lds_zo[s10]); } + } - rc_dec[31] = true; + /* Mantissa, MSB first */ + uint a = 1; + int i = int(e) - 1; + if (expectEXT(i >= 9, false)) { + /* Same for state 31 */ + uint s31 = SBYTE(31); + bool bb; + while (true) { + bb = rl_get(s31); + a = (a << 1) | uint(bb); + if (i == 9) + break; + s31 = zo_next(s31, bb); + i--; + } + i = 8; + sym_used |= 1u << 31; + sym_bits |= uint(bb) << 31; + if (lane == 7) { + dw = bitfieldInsert(dw, s31, 24, 8); + zt.w = uint(lds_zo[s31]); + } + } + + if (i == 8) { + bool bb = rl_get(SBYTE(30)); + a = (a << 1) | uint(bb); + sym_used |= 1u << 30; + sym_bits |= uint(bb) << 30; + i = 7; } - e += 20; - for (; e >= 22; e--) { - a <<= 1; - a |= int(READ(e)); - rc_dec[e] = true; + if (i >= 0) { + /* States 22..29, top one first */ + uint64_t m = (uint64_t(sb.y) >> 16) | (uint64_t(sb.z) << 16) | (uint64_t(sb.w) << 48); + m <<= 8*(7 - i); + do { + bool bb = rl_get(uint(m >> 56)); + m <<= 8; + a = (a << 1) | uint(bb); + sym_used |= 1u << (22 + i); + sym_bits |= uint(bb) << (22 + i); + } while (--i >= 0); } - return READ(sym_e) ? -a : a; + uint sidx = 11 + min(e, 10u); + bool sg = rl_get(sgs); + sym_used |= 1u << sidx; + sym_bits |= uint(sg) << sidx; + + return sg ? -int(a) : int(a); } void decode_line_pcm(ivec2 sp, int w, int y, int p) { - if (gl_LocalInvocationID.x > 0) - return; - #ifndef RGB if (p > 0 && p < 3) { w = ceil_rshift(w, chroma_shift.x); @@ -111,14 +308,18 @@ void decode_line_pcm(ivec2 sp, int w, int y, int p) uint v = 0; for (uint i = (rct_offset >> 1); i > 0; i >>= 1) - v |= get_rac_equi() ? i : 0; + v |= rl_get_equi() ? i : 0; - imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v)); + if (gl_SubgroupInvocationID == 0) + imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v)); } } +/* All samples are handled as TYPE, like the C decoder does */ +#define SX(v) int(TYPE(v)) + void decode_line(ivec2 sp, int w, - int y, int p, int bits, uint state_off, + int y, int p, int bits, uint gbase, uint lbase, uint8_t quant_table_idx, int run_index) { #ifndef RGB @@ -128,35 +329,141 @@ void decode_line(ivec2 sp, int w, } #endif - linecache_load(dec[p], sp, y, 0); + const uint lane = gl_SubgroupInvocationID; + const uint qi = uint(quant_table_idx); + const bool ext = has_extend_lookup && extend_lookup[quant_table_idx]; - for (int x = 0; x < w; x++) { - ivec2 pr = get_pred(dec[p], sp, ivec2(x, y), 0, w, - quant_table_idx, extend_lookup[quant_table_idx]); + uvec4 q0t = qt_load(qi, 0); + uvec4 q3t = uvec4(0); + if (ext) + q3t = qt_load(qi, 3); - uint rc_off = state_off + CONTEXT_SIZE*abs(pr[0]) + gl_LocalInvocationID.x; + int L = 0; +#ifdef RGB + L = SX(imageLoad(dec[p], sp + LADDR(ivec2(0, y - 1))).x); +#else + if (y > 0) + L = SX(imageLoad(dec[p], sp + ivec2(0, y - 1)).x); +#endif + int LL = 0; - rc_dec[gl_LocalInvocationID.x] = false; - rc_state[gl_LocalInvocationID.x] = slice_rc_state[rc_off]; - barrier(); + for (int x0 = 0; x0 < w; x0 += 32) { + /* Everything that only depends on previous lines is computed for + * the next 32 samples in parallel. */ + rl_window(); - if (gl_LocalInvocationID.x == 0) { - int diff = get_isymbol(); - if (pr[0] < 0) + int x = min(x0 + int(lane), w - 1); + int tl = 0, t = 0, tr = 0, tt = 0; + ivec2 tl_off = x == 0 ? ivec2(0, -2) : ivec2(x - 1, -1); + int tr_x = x + min(1, w - x - 1); +#ifdef RGB + tl = SX(imageLoad(dec[p], sp + LADDR(ivec2(0, y) + tl_off)).x); + t = SX(imageLoad(dec[p], sp + LADDR(ivec2(x, y - 1))).x); + tr = SX(imageLoad(dec[p], sp + LADDR(ivec2(tr_x, y - 1))).x); + if (ext) + tt = SX(imageLoad(dec[p], sp + LADDR(ivec2(x, rgb_linecache != 2 ? y - 2 : y))).x); +#else + if (y > 0) { + if (!(x == 0 && y == 1)) + tl = SX(imageLoad(dec[p], sp + ivec2(0, y) + tl_off).x); + t = SX(imageLoad(dec[p], sp + ivec2(x, y - 1)).x); + tr = SX(imageLoad(dec[p], sp + ivec2(tr_x, y - 1)).x); + } + if (ext && y > 1) + tt = SX(imageLoad(dec[p], sp + ivec2(x, y - 2)).x); +#endif + int kv = quant_table[qi][1][(tl - t) & MAX_QUANT_TABLE_MASK] + + quant_table[qi][2][(t - tr) & MAX_QUANT_TABLE_MASK]; + if (ext) + kv += quant_table[qi][4][(tt - t) & MAX_QUANT_TABLE_MASK]; + + int n = min(32, w - x0); + uint outv = 0; + + /* Fetched one sample ahead, off the critical path */ + int nT = subgroupBroadcast(t, 0); + int nTL = subgroupBroadcast(tl, 0); + int nK = subgroupBroadcast(kv, 0); + + for (int i = 0; i < n; i++) { + int T = nT, TL = nTL; + int ctx = nK + qt_get(q0t, uint(L - TL) & MAX_QUANT_TABLE_MASK); + if (ext) + ctx += qt_get(q3t, uint(LL - L) & MAX_QUANT_TABLE_MASK); + + int ni = min(i + 1, n - 1); + nT = subgroupBroadcast(t, ni); + nTL = subgroupBroadcast(tl, ni); + nK = subgroupBroadcast(kv, ni); + + int pred = mid_pred(L, L + T - TL, T); + + /* Invocation k < 8 owns states 4k..4k+3 of the context: it + * loads them, adapts them, and stores them back. */ + uint ci = (CONTEXT_SIZE >> 2)*uint(abs(ctx)) + min(lane, 7u); + uint dw; + if (use_lds) + dw = lds_state[(lbase >> 2) + ci]; + else + dw = slice_rc_state32[(gbase >> 2) + ci]; + + uvec4 sa = uvec4(subgroupBroadcast(dw, 0), subgroupBroadcast(dw, 1), + subgroupBroadcast(dw, 2), subgroupBroadcast(dw, 3)); + uvec4 sb = uvec4(subgroupBroadcast(dw, 4), subgroupBroadcast(dw, 5), + subgroupBroadcast(dw, 6), subgroupBroadcast(dw, 7)); + + /* Both possible adaptations of each state */ + uvec4 zt = uvec4(lds_zo[bitfieldExtract(dw, 0, 8)], + lds_zo[bitfieldExtract(dw, 8, 8)], + lds_zo[bitfieldExtract(dw, 16, 8)], + lds_zo[bitfieldExtract(dw, 24, 8)]); + + int diff = rl_get_isymbol(sa, sb, dw, zt); + if (ctx < 0) diff = -diff; + int v = zero_extend(pred + diff, bits); + + uint sh = 4*min(lane, 7u); + uint used4 = (sym_used >> sh) & 0xFu; + uint bits4 = (sym_bits >> sh) & 0xFu; + [[unroll]] + for (int k = 0; k < 4; k++) { + if (((used4 >> k) & 1u) != 0) + dw = bitfieldInsert(dw, ((bits4 >> k) & 1u) != 0 ? (zt[k] >> 8) : zt[k], 8*k, 8); + } - uint v = zero_extend(pr[1] + diff, bits); - imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v)); - linecache_next(TYPE(v)); + if (lane < 8) { + if (use_lds) + lds_state[(lbase >> 2) + ci] = dw; + else + slice_rc_state32[(gbase >> 2) + ci] = dw; + } + + outv = lane == i ? uint(v) : outv; + LL = L; + L = SX(v); } - /* Image write now visible to other invocs */ - barrier(); - if (rc_dec[gl_LocalInvocationID.x]) - slice_rc_state[rc_off] = - zero_one_state[rc_state[gl_LocalInvocationID.x] + - (rc_data[gl_LocalInvocationID.x] ? 256 : 0)]; + if (lane < n) + imageStore(dec[p], sp + LADDR(ivec2(x0 + int(lane), y)), uvec4(outv)); } + + /* Make the line visible for the next one */ + memoryBarrierImage(); + barrier(); +} + +/* Copy the probability states between VRAM and the on-chip copy */ +void lds_state_xfer(uint lbase, uint gbase, uint len, bool store) +{ + barrier(); + for (uint i = gl_LocalInvocationID.x; i < (len >> 2); i += gl_WorkGroupSize.x) { + if (store) + slice_rc_state32[(gbase >> 2) + i] = lds_state[(lbase >> 2) + i]; + else + lds_state[(lbase >> 2) + i] = slice_rc_state32[(gbase >> 2) + i]; + } + barrier(); } #else /* GOLOMB */ @@ -178,7 +485,7 @@ void golomb_init(void) } void decode_line(ivec2 sp, int w, - int y, int p, int bits, uint state_off, + int y, int p, int bits, uint state_off, uint lbase, uint8_t quant_table_idx, inout int run_index) { #ifndef RGB @@ -415,9 +722,31 @@ void decode_slice(in SliceContext sc, uint slice_idx) uvec4(0, 1, 1, 2))*plane_state_size; #endif + /* Offset of each plane's states within the on-chip copy */ + uint slice_base = slice_idx*codec_planes*plane_state_size; + u32vec4 lds_off = u32vec4(0); + #ifdef GOLOMB slice_state_off >>= 3; // division by VLC_STATE_SIZE golomb_init(); +#else + /* Only the contexts of the quantization tables used by this slice need + * to be kept on-chip */ + uint nb_ctx = 0; + for (int i = 0; i < codec_planes; i++) + nb_ctx = max(nb_ctx, uint(context_count[sc.quant_table_idx[i]])); + uint lds_stride = CONTEXT_SIZE*nb_ctx; +#ifdef RGB + /* All planes are coded line-interleaved: keep all of their states */ + use_lds = lds_state_bytes > 0 && (codec_planes*lds_stride) <= lds_state_bytes; + lds_off = ((slice_state_off - slice_base) / plane_state_size)*lds_stride; + if (use_lds) + for (int i = 0; i < codec_planes; i++) + lds_state_xfer(i*lds_stride, slice_base + i*plane_state_size, + lds_stride, false); +#else + use_lds = lds_state_bytes > 0 && lds_stride <= lds_state_bytes; +#endif #endif #ifdef BAYER @@ -425,7 +754,7 @@ void decode_slice(in SliceContext sc, uint slice_idx) for (int y = 0; y < bayer_h; y++) { for (int p = 0; p < 4; p++) decode_line(sp, w, y, p, bits[p], - slice_state_off[p], quant_table_idx[p], run_index); + slice_state_off[p], lds_off[p], quant_table_idx[p], run_index); writeout_bayer(slice_idx, sc, sp, w, y); } @@ -434,7 +763,7 @@ void decode_slice(in SliceContext sc, uint slice_idx) for (int y = 0; y < sc.slice_dim.y; y++) { for (int p = 0; p < color_planes; p++) decode_line(sp, w, y, p, bits[p], - slice_state_off[p], quant_table_idx[p], run_index); + slice_state_off[p], lds_off[p], quant_table_idx[p], run_index); writeout_rgb(slice_idx, sc, sp, w, y, true); } @@ -444,28 +773,59 @@ void decode_slice(in SliceContext sc, uint slice_idx) if (p > 0 && p < 3) h = ceil_rshift(h, chroma_shift.y); +#ifndef GOLOMB + /* Planes are coded one after another: keep only the current one. + * The two chroma planes share their states. */ + if (use_lds && p != 2) + lds_state_xfer(0, slice_state_off[p], lds_stride, false); +#endif + int run_index = 0; for (int y = 0; y < h; y++) decode_line(sp, w, y, p, bits[p], - slice_state_off[p], quant_table_idx[p], run_index); + slice_state_off[p], lds_off[p], quant_table_idx[p], run_index); + +#ifndef GOLOMB + if (use_lds && (p != 1 || (p + 1) >= planes)) + lds_state_xfer(0, slice_state_off[p], lds_stride, true); +#endif } #endif + +#if defined(RGB) && !defined(GOLOMB) + if (use_lds) + for (int i = 0; i < codec_planes; i++) + lds_state_xfer(i*lds_stride, slice_base + i*plane_state_size, + lds_stride, true); +#endif } void main(void) { uint slice_idx = gl_WorkGroupID.y*gl_NumWorkGroups.x + gl_WorkGroupID.x; +#ifndef GOLOMB + for (uint i = gl_LocalInvocationID.x; i < 256; i += gl_WorkGroupSize.x) + lds_zo[i] = uint16_t(zero_one_state[i]) | (uint16_t(zero_one_state[i + 256]) << 8); + rl_init(slice_ctx[slice_idx].c); + barrier(); +#else if (gl_LocalInvocationID.x == 0) rc = slice_ctx[slice_idx].c; barrier(); +#endif decode_slice(slice_ctx[slice_idx], slice_idx); if (gl_LocalInvocationID.x == 0) { uint overread = 0; +#ifndef GOLOMB + if (rl_off >= (rl_end + MAX_OVERREAD)) + overread = rl_off - rl_end; +#else if (rc.bs_off >= (rc.bs_end + MAX_OVERREAD)) overread = rc.bs_off - rc.bs_end; +#endif slice_status[2*slice_idx + 1] = overread; } } diff --git a/libavcodec/vulkan/ffv1_enc.comp.glsl b/libavcodec/vulkan/ffv1_enc.comp.glsl index f510266d7b..2a91e90e61 100644 --- a/libavcodec/vulkan/ffv1_enc.comp.glsl +++ b/libavcodec/vulkan/ffv1_enc.comp.glsl @@ -50,29 +50,191 @@ layout (set = 1, binding = 5, scalar) readonly buffer fltmap_buf { #ifndef GOLOMB +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_ballot : require + layout (set = 1, binding = 2, scalar) buffer slice_state_buf { uint8_t slice_rc_state[]; }; +layout (set = 1, binding = 2, scalar) buffer slice_state_buf32 { + uint32_t slice_rc_state32[]; +}; + +/* Size of the on-chip copy of the probability states, in bytes. When the + * states do not fit, they are accessed from slice_rc_state instead. */ +layout (constant_id = 19) const uint lds_state_bytes = 0; + +shared uint lds_state[lds_state_bytes > 0 ? (lds_state_bytes >> 2) : 1]; +/* Both possible next states of each state, as zero | (one << 8) */ +shared uint16_t lds_zo[256]; + +bool use_lds; + +uint zo_next(uint state, bool bit) +{ + uint z = uint(lds_zo[state]); + return bit ? (z >> 8) : (z & 0xFFu); +} + +/* + * Register-resident range encoder. + * The whole workgroup is a single subgroup, and every invocation performs + * the exact same computation on the same data, so all of this is + * dynamically uniform. Only the first invocation writes bytes out. + */ +uint re_low; +uint re_range; +uint re_off; /* Offset of the next byte to write */ +uint re_oc; /* Number of outstanding 0xFF bytes */ +uint re_ob; /* Outstanding byte */ -#define WRITE(idx, val) put_rac(rc_state[idx], val) -void put_symbol(int v) +void re_init(void) { - bool is_nil = (v == 0); - WRITE(0, is_nil); - if (is_nil) + re_low = rc.low; + re_range = rc.range; + re_off = rc.bs_off; + re_oc = uint(rc.outstanding_count); + re_ob = uint(rc.outstanding_byte) & 0xFFu; +} + +void re_sync(void) +{ + if (gl_SubgroupInvocationID == 0) { + rc.low = re_low; + rc.range = re_range; + rc.bs_off = re_off; + rc.outstanding_count = uint16_t(re_oc); + rc.outstanding_byte = int16_t(re_ob); + } + barrier(); +} + +void re_byte(uint v) +{ + if (gl_SubgroupInvocationID == 0) + slice_data[re_off].v = uint8_t(v); + re_off++; +} + +void re_renorm(void) +{ + uint low = re_low; + re_range <<= 8; + re_low = (low & 0xFFu) << 8; + if (low > 0xFF00 && low < 0x10000) { + re_oc++; return; + } - int a = abs(v); - int e = findMSB(a); + uint carry = uint(low > 0xFF00); + re_byte(re_ob + carry); + for (; re_oc > 0; re_oc--) + re_byte(carry - 1); + re_ob = (low >> 8) & 0xFFu; +} + +void re_put(uint state, bool bit) +{ + uint range1 = (re_range * state) >> 8; + uint ranged = re_range - range1; + re_low += bit ? ranged : 0; + re_range = bit ? range1 : ranged; + if (expectEXT(re_range < 0x100, false)) + re_renorm(); +} - for (int i = 0; i < e; i++) - WRITE(1 + min(i, 9), true); - WRITE(1 + min(e, 9), false); +/* The 32 probability states of the current context: dword k holds states + * 4k..4k+3. Every coded decision is known up front, so the masks of the + * states used, and of the values they coded, are computed directly, and + * all states get adapted in parallel after the symbol has been coded. */ +#define SBYTE(idx) bitfieldExtract((idx) < 16 ? sa[((idx) >> 2) & 3] : sb[((idx) >> 2) & 3], \ + ((idx) & 3)*8, 8) +uint sym_used; +uint sym_bits; - for (int i = e - 1; i >= 0; i--) - WRITE(22 + min(i, 9), bool(bitfieldExtract(a, i, 1))); +void re_put_symbol(uvec4 sa, uvec4 sb, int v, inout uint dw, inout uvec4 zt) +{ + const uint lane = gl_SubgroupInvocationID; - WRITE(22 - 11 + min(e, 10), v < 0); + bool nil = v == 0; + re_put(SBYTE(0), nil); + sym_used = 1u; + sym_bits = uint(nil); + if (nil) + return; + + uint a = uint(abs(v)); + int e = findMSB(a); + uint me = uint(min(e, 9)); + + /* Exponent: e ones, then a zero, with state 1 + min(k, 9) */ + uint64_t x = (uint64_t(sa.x) >> 8) | (uint64_t(sa.y) << 24) | (uint64_t(sa.z) << 56); + for (int k = 0; k < min(e, 8); k++) { + re_put(uint(x) & 0xFFu, true); + x >>= 8; + } + if (e < 8) { + re_put(uint(x) & 0xFFu, false); + } else { + re_put(bitfieldExtract(sa.z, 8, 8), e > 8); + if (e > 8) { + /* State 10 gets reused, and has to be adapted in between. + * The last write gets adapted along with the rest. */ + uint s10 = bitfieldExtract(sa.z, 16, 8); + for (int k = 9; k < e; k++) { + re_put(s10, true); + s10 = zo_next(s10, true); + } + re_put(s10, false); + if (lane == 2) { + dw = bitfieldInsert(dw, s10, 16, 8); + zt.z = uint(lds_zo[s10]); + } + } + } + sym_used |= ((2u << me) - 1u) << 1; + sym_bits |= ((1u << me) - 1u) << 1; + + /* Mantissa, MSB first, with state 22 + min(i, 9) */ + if (e >= 10) { + /* Same for state 31 */ + uint s31 = SBYTE(31); + for (int i = e - 1; i > 9; i--) { + bool bb = bitfieldExtract(a, i, 1) != 0; + re_put(s31, bb); + s31 = zo_next(s31, bb); + } + re_put(s31, bitfieldExtract(a, 9, 1) != 0); + sym_bits |= bitfieldExtract(a, 9, 1) << 31; + if (lane == 7) { + dw = bitfieldInsert(dw, s31, 24, 8); + zt.w = uint(lds_zo[s31]); + } + } + if (e >= 9) + re_put(SBYTE(30), bitfieldExtract(a, 8, 1) != 0); + int i = min(e, 8) - 1; + if (i >= 0) { + /* States 22..29, top one first */ + uint64_t m = (uint64_t(sb.y) >> 16) | (uint64_t(sb.z) << 16) | (uint64_t(sb.w) << 48); + m <<= 8*(7 - i); + do { + re_put(uint(m >> 56), bitfieldExtract(a, i, 1) != 0); + m <<= 8; + } while (--i >= 0); + } + sym_used |= ((1u << min(e, 10)) - 1u) << 22; + sym_bits |= (a & ((1u << me) - 1u)) << 22; + + /* Sign, with state 11 + min(e, 10) */ + uint se = uint(min(e, 10)); + uint64_t s64 = uint64_t(sa.z >> 24) | (uint64_t(sa.w) << 8) | (uint64_t(sb.x) << 40); + uint sgs = se < 8 ? (uint(s64 >> (se << 3)) & 0xFFu) : + se == 8 ? (sb.x >> 24) : + se == 9 ? (sb.y & 0xFFu) : bitfieldExtract(sb.y, 8, 8); + re_put(sgs, v < 0); + sym_used |= 1u << (11 + se); + sym_bits |= uint(v < 0) << (11 + se); } void encode_line_pcm(in SliceContext sc, readonly uimage2D img, @@ -101,7 +263,16 @@ void encode_line_pcm(in SliceContext sc, readonly uimage2D img, } } -void encode_line(in SliceContext sc, readonly uimage2D img, uint state_off, +/* All samples are handled as TYPE, like the C encoder does */ +#define SX(v) int(TYPE(v)) + +uint state_load(uint gbase, uint lbase, uint ctx) +{ + uint ci = (CONTEXT_SIZE >> 2)*ctx + min(gl_SubgroupInvocationID, 7u); + return use_lds ? lds_state[(lbase >> 2) + ci] : slice_rc_state32[(gbase >> 2) + ci]; +} + +void encode_line(in SliceContext sc, readonly uimage2D img, uint gbase, uint lbase, ivec2 sp, int y, uint p, uint comp, int bits, uint8_t quant_table_idx, in int run_index) { @@ -120,32 +291,116 @@ void encode_line(in SliceContext sc, readonly uimage2D img, uint state_off, return; #endif - linecache_load(img, sp, y, comp); - - for (int x = 0; x < w; x++) { - ivec2 d = get_pred(img, sp, ivec2(x, y), comp, w, - quant_table_idx, extend_lookup[quant_table_idx]); - TYPE cur = TYPE(imageLoad(img, sp + LADDR(ivec2(x, y)))[comp]); - d[1] = int(cur) - d[1]; - - if (d[0] < 0) - d = -d; - - d[1] = fold(d[1], bits); + const uint lane = gl_SubgroupInvocationID; + const uint qi = uint(quant_table_idx); + const bool ext = has_extend_lookup && extend_lookup[quant_table_idx]; - uint rc_off = state_off + CONTEXT_SIZE*d[0] + gl_LocalInvocationID.x; + for (int x0 = 0; x0 < w; x0 += 32) { + /* Contexts and residuals of the next 32 samples, in parallel */ + int x = min(x0 + int(lane), w - 1); + int tl = 0, t = 0, tr = 0, tt = 0; + ivec2 tl_off = x == 0 ? ivec2(0, -2) : ivec2(x - 1, -1); + int tr_x = x + min(1, w - x - 1); +#ifdef RGB + tl = SX(imageLoad(img, sp + LADDR(ivec2(0, y) + tl_off))[comp]); + t = SX(imageLoad(img, sp + LADDR(ivec2(x, y - 1)))[comp]); + tr = SX(imageLoad(img, sp + LADDR(ivec2(tr_x, y - 1)))[comp]); + if (ext) + tt = SX(imageLoad(img, sp + LADDR(ivec2(x, rgb_linecache != 2 ? y - 2 : y)))[comp]); +#else + if (y > 0) { + if (!(x == 0 && y == 1)) + tl = SX(imageLoad(img, sp + ivec2(0, y) + tl_off)[comp]); + t = SX(imageLoad(img, sp + ivec2(x, y - 1))[comp]); + tr = SX(imageLoad(img, sp + ivec2(tr_x, y - 1))[comp]); + } + if (ext && y > 1) + tt = SX(imageLoad(img, sp + ivec2(x, y - 2))[comp]); +#endif + int top0 = 0; + if (y > 0) + top0 = SX(imageLoad(img, sp + LADDR(ivec2(0, y - 1)))[comp]); + int L = x >= 1 ? SX(imageLoad(img, sp + LADDR(ivec2(x - 1, y)))[comp]) : top0; + int LL = x >= 2 ? SX(imageLoad(img, sp + LADDR(ivec2(x - 2, y)))[comp]) : + x == 1 ? top0 : 0; + int cur = SX(imageLoad(img, sp + LADDR(ivec2(x, y)))[comp]); + + int ctx = quant_table[qi][0][(L - tl) & MAX_QUANT_TABLE_MASK] + + quant_table[qi][1][(tl - t) & MAX_QUANT_TABLE_MASK] + + quant_table[qi][2][(t - tr) & MAX_QUANT_TABLE_MASK]; + if (ext) + ctx += quant_table[qi][3][(LL - L) & MAX_QUANT_TABLE_MASK] + + quant_table[qi][4][(tt - t) & MAX_QUANT_TABLE_MASK]; + + int res = cur - mid_pred(L, L + t - tl, t); + if (ctx < 0) { + ctx = -ctx; + res = -res; + } + res = fold(res, bits); + + /* Code them one by one. The states of the next sample are fetched + * while coding the current one, and forwarded when the context is + * the same. */ + int n = min(32, w - x0); + uint nc = uint(subgroupBroadcast(ctx, 0)); + uint ndw = state_load(gbase, lbase, nc); + uint pc = ~0u; + uint pdw = 0; + for (int i = 0; i < n; i++) { + uint c = nc; + uint dw = c == pc ? pdw : ndw; + int v = subgroupBroadcast(res, i); + if (i + 1 < n) { + nc = uint(subgroupBroadcast(ctx, i + 1)); + ndw = state_load(gbase, lbase, nc); + } - rc_state[gl_LocalInvocationID.x] = slice_rc_state[rc_off]; - barrier(); + uvec4 sa = uvec4(subgroupBroadcast(dw, 0), subgroupBroadcast(dw, 1), + subgroupBroadcast(dw, 2), subgroupBroadcast(dw, 3)); + uvec4 sb = uvec4(subgroupBroadcast(dw, 4), subgroupBroadcast(dw, 5), + subgroupBroadcast(dw, 6), subgroupBroadcast(dw, 7)); + uvec4 zt = uvec4(lds_zo[bitfieldExtract(dw, 0, 8)], + lds_zo[bitfieldExtract(dw, 8, 8)], + lds_zo[bitfieldExtract(dw, 16, 8)], + lds_zo[bitfieldExtract(dw, 24, 8)]); + + re_put_symbol(sa, sb, v, dw, zt); + + /* Invocation k < 8 owns states 4k..4k+3 of the context */ + uint sh = 4*min(lane, 7u); + uint used4 = (sym_used >> sh) & 0xFu; + uint bits4 = (sym_bits >> sh) & 0xFu; + [[unroll]] + for (int k = 0; k < 4; k++) { + if (((used4 >> k) & 1u) != 0) + dw = bitfieldInsert(dw, ((bits4 >> k) & 1u) != 0 ? (zt[k] >> 8) : zt[k], 8*k, 8); + } - if (gl_LocalInvocationID.x == 0) { - put_symbol(d[1]); - linecache_next(cur); + if (lane < 8) { + uint ci = (CONTEXT_SIZE >> 2)*c + lane; + if (use_lds) + lds_state[(lbase >> 2) + ci] = dw; + else + slice_rc_state32[(gbase >> 2) + ci] = dw; + } + pc = c; + pdw = dw; } + } +} - barrier(); - slice_rc_state[rc_off] = rc_state[gl_LocalInvocationID.x]; +/* Copy the probability states between VRAM and the on-chip copy */ +void lds_state_xfer(uint lbase, uint gbase, uint len, bool store) +{ + barrier(); + for (uint i = gl_LocalInvocationID.x; i < (len >> 2); i += gl_WorkGroupSize.x) { + if (store) + slice_rc_state32[(gbase >> 2) + i] = lds_state[(lbase >> 2) + i]; + else + lds_state[(lbase >> 2) + i] = slice_rc_state32[(gbase >> 2) + i]; } + barrier(); } #else /* GOLOMB */ @@ -164,7 +419,7 @@ void init_golomb(void) slice_size_max - hdr_len); } -void encode_line(in SliceContext sc, readonly uimage2D img, uint state_off, +void encode_line(in SliceContext sc, readonly uimage2D img, uint state_off, uint lbase, ivec2 sp, int y, uint p, uint comp, int bits, uint8_t quant_table_idx, inout int run_index) { @@ -414,9 +669,23 @@ void encode_slice(in SliceContext sc, uint slice_idx) uvec4(0, 1, 1, 2))*plane_state_size; #endif + uint slice_base = slice_idx*codec_planes*plane_state_size; + u32vec4 lds_off = u32vec4(0); + #ifdef GOLOMB slice_state_off >>= 3; init_golomb(); +#else + re_init(); +#ifdef RGB + /* All planes are coded line-interleaved: keep all of their states */ + use_lds = lds_state_bytes > 0 && (codec_planes*plane_state_size) <= lds_state_bytes; + lds_off = slice_state_off - slice_base; + if (use_lds) + lds_state_xfer(0, slice_base, codec_planes*plane_state_size, false); +#else + use_lds = lds_state_bytes > 0 && plane_state_size <= lds_state_bytes; +#endif #endif #ifdef BAYER @@ -425,7 +694,7 @@ void encode_slice(in SliceContext sc, uint slice_idx) preload_bayer(sc, sp, bayer_w, y, true); for (uint c = 0; c < 4; c++) - encode_line(sc, tmp, slice_state_off[c], + encode_line(sc, tmp, slice_state_off[c], lds_off[c], sp, y, 0, c, bits[c], U8(context_model), run_index); } @@ -435,7 +704,7 @@ void encode_slice(in SliceContext sc, uint slice_idx) preload_rgb(slice_idx, sc, sp, sc.slice_dim.x, y, true); for (uint c = 0; c < color_planes; c++) - encode_line(sc, tmp, slice_state_off[c], + encode_line(sc, tmp, slice_state_off[c], lds_off[c], sp, y, 0, rgb_plane_order[c], bits[c], U8(context_model), run_index); } @@ -450,11 +719,31 @@ void encode_slice(in SliceContext sc, uint slice_idx) uint p = min(c, planes - 1); uint comp = c - p; +#ifndef GOLOMB + /* Planes are coded one after another: keep only the current one. + * The two chroma planes share their states. */ + if (use_lds && c != 2) + lds_state_xfer(0, slice_state_off[c], plane_state_size, false); +#endif + for (int y = 0; y < h; y++) - encode_line(sc, src[p], slice_state_off[c], sp, y, p, + encode_line(sc, src[p], slice_state_off[c], 0, sp, y, p, comp, bits[c], U8(context_model), run_index); + +#ifndef GOLOMB + if (use_lds && (c != 1 || (c + 1) >= color_planes)) + lds_state_xfer(0, slice_state_off[c], plane_state_size, true); +#endif } #endif + +#if defined(RGB) && !defined(GOLOMB) + if (use_lds) + lds_state_xfer(0, slice_base, codec_planes*plane_state_size, true); +#endif +#ifndef GOLOMB + re_sync(); +#endif } void finalize_slice(in uint slice_idx) @@ -501,6 +790,10 @@ void main(void) { uint slice_idx = gl_WorkGroupID.y*gl_NumWorkGroups.x + gl_WorkGroupID.x; +#ifndef GOLOMB + for (uint i = gl_LocalInvocationID.x; i < 256; i += gl_WorkGroupSize.x) + lds_zo[i] = uint16_t(zero_one_state[i]) | (uint16_t(zero_one_state[i + 256]) << 8); +#endif if (gl_LocalInvocationID.x == 0) rc = slice_ctx[slice_idx].c; barrier(); diff --git a/libavcodec/vulkan_ffv1.c b/libavcodec/vulkan_ffv1.c index 779f316c51..d02f10d6c2 100644 --- a/libavcodec/vulkan_ffv1.c +++ b/libavcodec/vulkan_ffv1.c @@ -692,8 +692,15 @@ static int init_decode_shader(FFV1Context *f, FFVulkanContext *s, int err; uint32_t wg_x = ac != AC_GOLOMB_RICE ? CONTEXT_SIZE : 1; + uint32_t sg_size = 0; + /* The range coder decoder needs its workgroup to be a single subgroup */ + if (ac != AC_GOLOMB_RICE && + (s->subgroup_props.requiredSubgroupSizeStages & VK_SHADER_STAGE_COMPUTE_BIT) && + s->subgroup_props.minSubgroupSize <= CONTEXT_SIZE && + s->subgroup_props.maxSubgroupSize >= CONTEXT_SIZE) + sg_size = CONTEXT_SIZE; ff_vk_shader_load(shd, VK_SHADER_STAGE_COMPUTE_BIT, sl, - (uint32_t []) { wg_x, 1, 1 }, 0); + (uint32_t []) { wg_x, 1, 1 }, sg_size); ff_vk_shader_add_push_const(shd, 0, sizeof(FFv1ShaderParams), VK_SHADER_STAGE_COMPUTE_BIT); @@ -897,7 +904,7 @@ static int vk_decode_ffv1_init(AVCodecContext *avctx) dctx = (AVHWFramesContext *)fv->intermediate_frames_ref->data; } - SPEC_LIST_CREATE(sl, 15, 15*sizeof(uint32_t)) + SPEC_LIST_CREATE(sl, 16, 16*sizeof(uint32_t)) ff_ffv1_vk_set_common_sl(avctx, f, sl, sw_format); if (RGB_LINECACHE != 2) @@ -912,6 +919,26 @@ static int vk_decode_ffv1_init(AVCodecContext *avctx) /* Reset shader */ RET(init_reset_shader(f, &ctx->s, &ctx->exec_pool, &fv->reset, sl, f->ac)); + /* Keep the probability states on-chip if they fit. Line-interleaved + * planes need all of their states at once. Each slice picks its own + * quantization tables, so size this for the largest set that fits, and + * let the shader decide per slice. */ + if (f->ac != AC_GOLOMB_RICE) { + uint32_t lds_bytes = 0; + uint32_t lds_max = ctx->s.props.properties.limits.maxComputeSharedMemorySize; + lds_max = lds_max > 2048 ? lds_max - 2048 : 0; + for (int i = 0; i < f->quant_table_count; i++) { + uint32_t len = CONTEXT_SIZE*f->context_count[i]; + if (is_rgb || f->bayer) + len *= f->plane_count; + if (len <= lds_max) + lds_bytes = FFMAX(lds_bytes, len); + } + SPEC_LIST_ADD(sl, 19, 32, lds_bytes); + av_log(avctx, AV_LOG_VERBOSE, "FFv1: up to %u bytes of on-chip probability states\n", + lds_bytes); + } + /* Decode shaders */ RET(init_decode_shader(f, &ctx->s, &ctx->exec_pool, &fv->decode, dctx, hwfc, sl, f->ac, is_rgb, f->bayer));