libavcodec/vulkan/ffv1_dec.comp.glsl - third_party/ffmpeg - Git at Google

 /*
  * FFv1 codec
  *
  * Copyright (c) 2024 Lynne <dev@lynne.ee>
  *
  * This file is part of FFmpeg.
  *
  * FFmpeg is free software; you can redistribute it and/or
  * modify it under the terms of the GNU Lesser General Public
  * License as published by the Free Software Foundation; either
  * version 2.1 of the License, or (at your option) any later version.
  *
  * FFmpeg is distributed in the hope that it will be useful,
  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  * Lesser General Public License for more details.
  *
  * You should have received a copy of the GNU Lesser General Public
  * License along with FFmpeg; if not, write to the Free Software
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  */

 #pragma shader_stage(compute)
 #extension GL_GOOGLE_include_directive : require

 #define DECODE
 #include "common.glsl"
 #include "ffv1_common.glsl"

 layout (set = 1, binding = 1, scalar) readonly buffer slice_offsets_buf {
     u32vec2 slice_offsets[];
 };
 layout (set = 1, binding = 2, scalar) writeonly buffer slice_status_buf {
     uint32_t slice_status[];
 };
 layout (set = 1, binding = 4) uniform uimage2D dec[];

 #ifndef GOLOMB

 layout (set = 1, binding = 3, scalar) buffer slice_state_buf {
     uint8_t slice_rc_state[];
 };

 #define READ(idx) get_rac_state(idx)
 shared int sym_e;
 shared bool rc_dec[CONTEXT_SIZE];
 int get_isymbol(void)
 {
     sym_e = 0;
     rc_dec[0] = true;
     if (READ(0))
         return 0;

     int e = 1;
     for (; e < 11; e++) {
         rc_dec[e] = true;
         if (!READ(e))
             break;
     }

     int a = 1;
     sym_e = e + 10;
     rc_dec[sym_e] = true;

     if (bits > 10 && e == 11) {
         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));
         }

         rc_dec[31] = true;
     }

     e += 20;
     for (; e >= 22; e--) {
         a <<= 1;
         a |= int(READ(e));
         rc_dec[e] = true;
     }

     return READ(sym_e) ? -a : 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);
         sp >>= chroma_shift;
     }
 #endif

     for (int x = 0; x < w; x++) {
         uint v = 0;

         for (uint i = (rct_offset >> 1); i > 0; i >>= 1)
             v |= get_rac_equi() ? i : 0;

         imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v));
     }
 }

 void decode_line(ivec2 sp, int w,
                  int y, int p, uint state_off,
                  uint8_t quant_table_idx, int run_index)
 {
 #ifndef RGB
     if (p > 0 && p < 3) {
         w = ceil_rshift(w, chroma_shift.x);
         sp >>= chroma_shift;
     }
 #endif

     linecache_load(dec[p], sp, y, 0);

     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]);

         uint rc_off = state_off + CONTEXT_SIZE*abs(pr[0]) + gl_LocalInvocationID.x;

         rc_dec[gl_LocalInvocationID.x] = false;
         rc_state[gl_LocalInvocationID.x] = slice_rc_state[rc_off];
         barrier();

         if (gl_LocalInvocationID.x == 0) {
             int diff = get_isymbol();
             if (pr[0] < 0)
                 diff = -diff;

             uint v = zero_extend(pr[1] + diff, bits);
             imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v));
             linecache_next(TYPE(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)];
     }
 }

 #else /* GOLOMB */

 layout (set = 1, binding = 3, scalar) buffer slice_state_buf {
     VlcState slice_vlc_state[];
 };

 GetBitContext gb;

 void golomb_init(void)
 {
     if (version == 3 && micro_version > 1 || version > 3)
         get_rac_internal(rc.range * 129 >> 8);

     uint64_t ac_byte_count = rc.bs_off - rc.bs_start - 1;
     init_get_bits(gb, u8buf(rc.bs_start + ac_byte_count),
                   int(rc.bs_end - rc.bs_start - ac_byte_count));
 }

 void decode_line(ivec2 sp, int w,
                  int y, int p, uint state_off,
                  uint8_t quant_table_idx, inout int run_index)
 {
 #ifndef RGB
     if (p > 0 && p < 3) {
         w = ceil_rshift(w, chroma_shift.x);
         sp >>= chroma_shift;
     }
 #endif

     linecache_load(dec[p], sp, y, 0);

     int run_count = 0;
     int run_mode  = 0;

     for (int x = 0; x < w; x++) {
         ivec2 pos = sp + ivec2(x, y);
         int diff;
         ivec2 pr = get_pred(dec[p], sp, ivec2(x, y), 0, w,
                             quant_table_idx, extend_lookup[quant_table_idx]);

         uint vlc_off = state_off + abs(pr[0]);

         if (pr[0] == 0 && run_mode == 0)
             run_mode = 1;

         if (run_mode != 0) {
             if (run_count == 0 && run_mode == 1) {
                 int tmp_idx = int(log2_run[run_index]);
                 if (get_bit(gb)) {
                     run_count = 1 << tmp_idx;
                     if (x + run_count <= w)
                         run_index++;
                 } else {
                     if (tmp_idx != 0) {
                         run_count = int(get_bits(gb, tmp_idx));
                     } else
                         run_count = 0;

                     if (run_index != 0)
                         run_index--;
                     run_mode = 2;
                 }
             }

             run_count--;
             if (run_count < 0) {
                 run_mode  = 0;
                 run_count = 0;
                 diff = read_vlc_symbol(gb, slice_vlc_state[vlc_off], bits);
                 if (diff >= 0)
                     diff++;
             } else {
                 diff = 0;
             }
         } else {
             diff = read_vlc_symbol(gb, slice_vlc_state[vlc_off], bits);
         }

         if (pr[0] < 0)
             diff = -diff;

         uint v = zero_extend(pr[1] + diff, bits);
         imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v));
         linecache_next(TYPE(v));
     }
 }
 #endif

 #ifdef RGB
 ivec4 transform_sample(ivec4 pix, ivec2 rct_coef)
 {
     pix.b -= rct_offset;
     pix.r -= rct_offset;
     pix.g -= (pix.b*rct_coef.g + pix.r*rct_coef.r) >> 2;
     pix.b += pix.g;
     pix.r += pix.g;
     return ivec4(pix[fmt_lut[0]], pix[fmt_lut[1]],
                  pix[fmt_lut[2]], pix[fmt_lut[3]]);
 }

 void writeout_rgb(in SliceContext sc, ivec2 sp, int w, int y, bool apply_rct)
 {
     memoryBarrierImage();
     barrier();

     for (uint x = gl_LocalInvocationID.x; x < w; x += gl_WorkGroupSize.x) {
         ivec2 lpos = sp + LADDR(ivec2(x, y));
         ivec2 pos = sc.slice_pos + ivec2(x, y);

         ivec4 pix;
         pix.r = int(imageLoad(dec[2], lpos)[0]);
         pix.g = int(imageLoad(dec[0], lpos)[0]);
         pix.b = int(imageLoad(dec[1], lpos)[0]);
         if (transparency)
             pix.a = int(imageLoad(dec[3], lpos)[0]);

         if (apply_rct)
             pix = transform_sample(pix, sc.slice_rct_coef);
         else
             pix = ivec4(pix[fmt_lut[0]], pix[fmt_lut[1]],
                         pix[fmt_lut[2]], pix[fmt_lut[3]]);

         imageStore(dst[0], pos, pix);
         if (planar_rgb) {
             for (int i = 1; i < color_planes; i++)
                 imageStore(dst[i], pos, ivec4(pix[i]));
         }
     }
 }
 #endif

 void decode_slice(in SliceContext sc, uint slice_idx)
 {
     int w = sc.slice_dim.x;
     ivec2 sp = sc.slice_pos;

 #ifdef RGB
     sp.y = int(gl_WorkGroupID.y)*rgb_linecache;
 #endif

 #ifndef GOLOMB
     /* PCM coding */
     if (sc.slice_coding_mode == 1) {
 #ifdef RGB
         for (int y = 0; y < sc.slice_dim.y; y++) {
             for (int p = 0; p < color_planes; p++)
                 decode_line_pcm(sp, w, y, p);

             writeout_rgb(sc, sp, w, y, false);
         }
 #else
         for (int p = 0; p < planes; p++) {
             int h = sc.slice_dim.y;
             if (p > 0 && p < 3)
                 h = ceil_rshift(h, chroma_shift.y);

             for (int y = 0; y < h; y++)
                 decode_line_pcm(sp, w, y, p);
         }
 #endif
         return;
     }
 #endif

     u8vec4 quant_table_idx = sc.quant_table_idx.xyyz;
     u32vec4 slice_state_off = (slice_idx*codec_planes +
                                uvec4(0, 1, 1, 2))*plane_state_size;

 #ifdef GOLOMB
     slice_state_off >>= 3; // division by VLC_STATE_SIZE
     golomb_init();
 #endif

 #ifdef RGB
     int run_index = 0;
     for (int y = 0; y < sc.slice_dim.y; y++) {
         for (int p = 0; p < color_planes; p++)
             decode_line(sp, w, y, p,
                         slice_state_off[p], quant_table_idx[p], run_index);

         writeout_rgb(sc, sp, w, y, true);
     }
 #else
     for (int p = 0; p < planes; p++) {
         int h = sc.slice_dim.y;
         if (p > 0 && p < 3)
             h = ceil_rshift(h, chroma_shift.y);

         int run_index = 0;
         for (int y = 0; y < h; y++)
             decode_line(sp, w, y, p,
                         slice_state_off[p], quant_table_idx[p], run_index);
     }
 #endif
 }

 void main(void)
 {
     uint slice_idx = gl_WorkGroupID.y*gl_NumWorkGroups.x + gl_WorkGroupID.x;

     if (gl_LocalInvocationID.x == 0)
         rc = slice_ctx[slice_idx].c;
     barrier();

     decode_slice(slice_ctx[slice_idx], slice_idx);

     if (gl_LocalInvocationID.x == 0) {
         uint overread = 0;
         if (rc.bs_off >= (rc.bs_end + MAX_OVERREAD))
             overread = rc.bs_off - rc.bs_end;
         slice_status[2*slice_idx + 1] = overread;
     }
 }
	/*
	* FFv1 codec
	*
	* Copyright (c) 2024 Lynne <dev@lynne.ee>
	*
	* This file is part of FFmpeg.
	*
	* FFmpeg is free software; you can redistribute it and/or
	* modify it under the terms of the GNU Lesser General Public
	* License as published by the Free Software Foundation; either
	* version 2.1 of the License, or (at your option) any later version.
	*
	* FFmpeg is distributed in the hope that it will be useful,
	* but WITHOUT ANY WARRANTY; without even the implied warranty of
	* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
	* Lesser General Public License for more details.
	*
	* You should have received a copy of the GNU Lesser General Public
	* License along with FFmpeg; if not, write to the Free Software
	* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
	*/

	#pragma shader_stage(compute)
	#extension GL_GOOGLE_include_directive : require

	#define DECODE
	#include "common.glsl"
	#include "ffv1_common.glsl"

	layout (set = 1, binding = 1, scalar) readonly buffer slice_offsets_buf {
	u32vec2 slice_offsets[];
	};
	layout (set = 1, binding = 2, scalar) writeonly buffer slice_status_buf {
	uint32_t slice_status[];
	};
	layout (set = 1, binding = 4) uniform uimage2D dec[];

	#ifndef GOLOMB

	layout (set = 1, binding = 3, scalar) buffer slice_state_buf {
	uint8_t slice_rc_state[];
	};

	#define READ(idx) get_rac_state(idx)
	shared int sym_e;
	shared bool rc_dec[CONTEXT_SIZE];
	int get_isymbol(void)
	{
	sym_e = 0;
	rc_dec[0] = true;
	if (READ(0))
	return 0;

	int e = 1;
	for (; e < 11; e++) {
	rc_dec[e] = true;
	if (!READ(e))
	break;
	}

	int a = 1;
	sym_e = e + 10;
	rc_dec[sym_e] = true;

	if (bits > 10 && e == 11) {
	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));
	}

	rc_dec[31] = true;
	}

	e += 20;
	for (; e >= 22; e--) {
	a <<= 1;
	a \|= int(READ(e));
	rc_dec[e] = true;
	}

	return READ(sym_e) ? -a : 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);
	sp >>= chroma_shift;
	}
	#endif

	for (int x = 0; x < w; x++) {
	uint v = 0;

	for (uint i = (rct_offset >> 1); i > 0; i >>= 1)
	v \|= get_rac_equi() ? i : 0;

	imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v));
	}
	}

	void decode_line(ivec2 sp, int w,
	int y, int p, uint state_off,
	uint8_t quant_table_idx, int run_index)
	{
	#ifndef RGB
	if (p > 0 && p < 3) {
	w = ceil_rshift(w, chroma_shift.x);
	sp >>= chroma_shift;
	}
	#endif

	linecache_load(dec[p], sp, y, 0);

	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]);

	uint rc_off = state_off + CONTEXT_SIZE*abs(pr[0]) + gl_LocalInvocationID.x;

	rc_dec[gl_LocalInvocationID.x] = false;
	rc_state[gl_LocalInvocationID.x] = slice_rc_state[rc_off];
	barrier();

	if (gl_LocalInvocationID.x == 0) {
	int diff = get_isymbol();
	if (pr[0] < 0)
	diff = -diff;

	uint v = zero_extend(pr[1] + diff, bits);
	imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v));
	linecache_next(TYPE(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)];
	}
	}

	#else /* GOLOMB */

	layout (set = 1, binding = 3, scalar) buffer slice_state_buf {
	VlcState slice_vlc_state[];
	};

	GetBitContext gb;

	void golomb_init(void)
	{
	if (version == 3 && micro_version > 1 \|\| version > 3)
	get_rac_internal(rc.range * 129 >> 8);

	uint64_t ac_byte_count = rc.bs_off - rc.bs_start - 1;
	init_get_bits(gb, u8buf(rc.bs_start + ac_byte_count),
	int(rc.bs_end - rc.bs_start - ac_byte_count));
	}

	void decode_line(ivec2 sp, int w,
	int y, int p, uint state_off,
	uint8_t quant_table_idx, inout int run_index)
	{
	#ifndef RGB
	if (p > 0 && p < 3) {
	w = ceil_rshift(w, chroma_shift.x);
	sp >>= chroma_shift;
	}
	#endif

	linecache_load(dec[p], sp, y, 0);

	int run_count = 0;
	int run_mode = 0;

	for (int x = 0; x < w; x++) {
	ivec2 pos = sp + ivec2(x, y);
	int diff;
	ivec2 pr = get_pred(dec[p], sp, ivec2(x, y), 0, w,
	quant_table_idx, extend_lookup[quant_table_idx]);

	uint vlc_off = state_off + abs(pr[0]);

	if (pr[0] == 0 && run_mode == 0)
	run_mode = 1;

	if (run_mode != 0) {
	if (run_count == 0 && run_mode == 1) {
	int tmp_idx = int(log2_run[run_index]);
	if (get_bit(gb)) {
	run_count = 1 << tmp_idx;
	if (x + run_count <= w)
	run_index++;
	} else {
	if (tmp_idx != 0) {
	run_count = int(get_bits(gb, tmp_idx));
	} else
	run_count = 0;

	if (run_index != 0)
	run_index--;
	run_mode = 2;
	}
	}

	run_count--;
	if (run_count < 0) {
	run_mode = 0;
	run_count = 0;
	diff = read_vlc_symbol(gb, slice_vlc_state[vlc_off], bits);
	if (diff >= 0)
	diff++;
	} else {
	diff = 0;
	}
	} else {
	diff = read_vlc_symbol(gb, slice_vlc_state[vlc_off], bits);
	}

	if (pr[0] < 0)
	diff = -diff;

	uint v = zero_extend(pr[1] + diff, bits);
	imageStore(dec[p], sp + LADDR(ivec2(x, y)), uvec4(v));
	linecache_next(TYPE(v));
	}
	}
	#endif

	#ifdef RGB
	ivec4 transform_sample(ivec4 pix, ivec2 rct_coef)
	{
	pix.b -= rct_offset;
	pix.r -= rct_offset;
	pix.g -= (pix.brct_coef.g + pix.rrct_coef.r) >> 2;
	pix.b += pix.g;
	pix.r += pix.g;
	return ivec4(pix[fmt_lut[0]], pix[fmt_lut[1]],
	pix[fmt_lut[2]], pix[fmt_lut[3]]);
	}

	void writeout_rgb(in SliceContext sc, ivec2 sp, int w, int y, bool apply_rct)
	{
	memoryBarrierImage();
	barrier();

	for (uint x = gl_LocalInvocationID.x; x < w; x += gl_WorkGroupSize.x) {
	ivec2 lpos = sp + LADDR(ivec2(x, y));
	ivec2 pos = sc.slice_pos + ivec2(x, y);

	ivec4 pix;
	pix.r = int(imageLoad(dec[2], lpos)[0]);
	pix.g = int(imageLoad(dec[0], lpos)[0]);
	pix.b = int(imageLoad(dec[1], lpos)[0]);
	if (transparency)
	pix.a = int(imageLoad(dec[3], lpos)[0]);

	if (apply_rct)
	pix = transform_sample(pix, sc.slice_rct_coef);
	else
	pix = ivec4(pix[fmt_lut[0]], pix[fmt_lut[1]],
	pix[fmt_lut[2]], pix[fmt_lut[3]]);

	imageStore(dst[0], pos, pix);
	if (planar_rgb) {
	for (int i = 1; i < color_planes; i++)
	imageStore(dst[i], pos, ivec4(pix[i]));
	}
	}
	}
	#endif

	void decode_slice(in SliceContext sc, uint slice_idx)
	{
	int w = sc.slice_dim.x;
	ivec2 sp = sc.slice_pos;

	#ifdef RGB
	sp.y = int(gl_WorkGroupID.y)*rgb_linecache;
	#endif

	#ifndef GOLOMB
	/* PCM coding */
	if (sc.slice_coding_mode == 1) {
	#ifdef RGB
	for (int y = 0; y < sc.slice_dim.y; y++) {
	for (int p = 0; p < color_planes; p++)
	decode_line_pcm(sp, w, y, p);

	writeout_rgb(sc, sp, w, y, false);
	}
	#else
	for (int p = 0; p < planes; p++) {
	int h = sc.slice_dim.y;
	if (p > 0 && p < 3)
	h = ceil_rshift(h, chroma_shift.y);

	for (int y = 0; y < h; y++)
	decode_line_pcm(sp, w, y, p);
	}
	#endif
	return;
	}
	#endif

	u8vec4 quant_table_idx = sc.quant_table_idx.xyyz;
	u32vec4 slice_state_off = (slice_idx*codec_planes +
	uvec4(0, 1, 1, 2))*plane_state_size;

	#ifdef GOLOMB
	slice_state_off >>= 3; // division by VLC_STATE_SIZE
	golomb_init();
	#endif

	#ifdef RGB
	int run_index = 0;
	for (int y = 0; y < sc.slice_dim.y; y++) {
	for (int p = 0; p < color_planes; p++)
	decode_line(sp, w, y, p,
	slice_state_off[p], quant_table_idx[p], run_index);

	writeout_rgb(sc, sp, w, y, true);
	}
	#else
	for (int p = 0; p < planes; p++) {
	int h = sc.slice_dim.y;
	if (p > 0 && p < 3)
	h = ceil_rshift(h, chroma_shift.y);

	int run_index = 0;
	for (int y = 0; y < h; y++)
	decode_line(sp, w, y, p,
	slice_state_off[p], quant_table_idx[p], run_index);
	}
	#endif
	}

	void main(void)
	{
	uint slice_idx = gl_WorkGroupID.y*gl_NumWorkGroups.x + gl_WorkGroupID.x;

	if (gl_LocalInvocationID.x == 0)
	rc = slice_ctx[slice_idx].c;
	barrier();

	decode_slice(slice_ctx[slice_idx], slice_idx);

	if (gl_LocalInvocationID.x == 0) {
	uint overread = 0;
	if (rc.bs_off >= (rc.bs_end + MAX_OVERREAD))
	overread = rc.bs_off - rc.bs_end;
	slice_status[2*slice_idx + 1] = overread;
	}
	}