capcodec/tools/gen_tables.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
import re
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
JM = Path.home() / "tools" / "JM-master" / "source"
FF = Path.home() / "work" / "refsrc"
OUT = ROOT / "src" / "tables.tov"
# Tables wider than their values need: the quantisers are loaded as i32 vectors, and byte loads
# of the zigzag and deblocking tables cost instructions in their hot loops (callgrind).
ELEM = {
"QUANT4": "i32", "DEQUANT4": "i32", "QUANT8": "i32", "DEQUANT8": "i32",
"ZIGZAG4": "int", "ZIGZAG8": "int", "DEBLOCK_ALPHA": "int", "DEBLOCK_BETA": "int", "DEBLOCK_TC0": "int",
}
def strip_comments(text: str) -> str:
text = re.sub(r"/\*.*?\*/", " ", text, flags=re.S)
return re.sub(r"//[^\n]*", " ", text)
def parse_init(text: str, start: int):
i = text.index("{", start)
tokens = re.findall(r"\{|\}|-?\d+|CTX_UNUSED|,", text[i:])
pos = 0
def parse():
nonlocal pos
assert tokens[pos] == "{"
pos += 1
items = []
while tokens[pos] != "}":
t = tokens[pos]
if t == ",":
pos += 1
continue
if t == "{":
items.append(parse())
elif t == "CTX_UNUSED":
items.append([0, 0])
pos += 1
else:
items.append(int(t))
pos += 1
pos += 1
return items
return parse()
def shape_fill(v, shape):
if not shape:
return v if isinstance(v, int) else 0
out = []
v = v if isinstance(v, list) else []
for k in range(shape[0]):
out.append(shape_fill(v[k] if k < len(v) else 0, shape[1:]))
return out
def flat(v):
if isinstance(v, int):
return [v]
out = []
for x in v:
out += flat(x)
return out
def table(path: Path, decl: str, shape, after: str | None = None):
text = strip_comments(path.read_text(errors="replace"))
start = 0
if after:
start = text.index(after)
m = re.compile(re.escape(decl)).search(text, start)
if not m:
raise SystemExit(f"{decl} not found in {path}")
return flat(shape_fill(parse_init(text, m.end()), shape))
def narrowest(values: list[int]) -> str:
lo, hi = min(values), max(values)
for elem, a, b in (("u8", 0, 255), ("i8", -128, 127), ("u16", 0, 65535), ("i16", -32768, 32767)):
if a <= lo and hi <= b:
return elem
return "i32"
def emit(name: str, values: list[int], per_line: int = 32, elem: str | None = None) -> str:
elem = elem or narrowest(values)
lines = []
for i in range(0, len(values), per_line):
lines.append(" " + ", ".join(str(v) for v in values[i:i + per_line]) + ",")
return f"export const {name}: {elem}[] = [\n" + "\n".join(lines) + "\n]\n"
def main() -> int:
vlc = JM / "app" / "lencod" / "vlc.c"
t = {}
t["CABAC_INIT_I"] = table(FF / "h264_cabac.c", "cabac_context_init_I[1024][2] =", [1024, 2])
t["CABAC_INIT_P"] = table(FF / "h264_cabac.c", "cabac_context_init_PB[3][1024][2] =", [3, 1024, 2])
bia = JM / "app" / "lencod" / "biariencode.c"
t["CABAC_RANGE_LPS"] = table(bia, "rLPS_table_64x4[64][4]=", [64, 4])
t["CABAC_NEXT_MPS"] = table(bia, "AC_next_state_MPS_64[64] =", [64])
t["CABAC_NEXT_LPS"] = table(bia, "AC_next_state_LPS_64[64] =", [64])
t["COEFF_TOKEN_LEN"] = table(vlc, "lentab[3][4][17] =", [3, 4, 17], "writeSyntaxElement_NumCoeffTrailingOnes(")
t["COEFF_TOKEN_CODE"] = table(vlc, "codtab[3][4][17] =", [3, 4, 17], "writeSyntaxElement_NumCoeffTrailingOnes(")
t["COEFF_TOKEN_DC_LEN"] = table(vlc, "lentab[3][4][17] =", [3, 4, 17],
"writeSyntaxElement_NumCoeffTrailingOnesChromaDC(")[:68]
t["COEFF_TOKEN_DC_CODE"] = table(vlc, "codtab[3][4][17] =", [3, 4, 17],
"writeSyntaxElement_NumCoeffTrailingOnesChromaDC(")[:68]
t["TOTAL_ZEROS_LEN"] = table(vlc, "lentab[TOTRUN_NUM][16] =", [15, 16], "writeSyntaxElement_TotalZeros(")
t["TOTAL_ZEROS_CODE"] = table(vlc, "codtab[TOTRUN_NUM][16] =", [15, 16], "writeSyntaxElement_TotalZeros(")
t["TOTAL_ZEROS_DC_LEN"] = table(vlc, "lentab[3][TOTRUN_NUM][16] =", [3, 15, 16],
"writeSyntaxElement_TotalZerosChromaDC(")[:240]
t["TOTAL_ZEROS_DC_CODE"] = table(vlc, "codtab[3][TOTRUN_NUM][16] =", [3, 15, 16],
"writeSyntaxElement_TotalZerosChromaDC(")[:240]
t["RUN_BEFORE_LEN"] = table(vlc, "lentab[TOTRUN_NUM][16] =", [15, 16], "writeSyntaxElement_Run(")
t["RUN_BEFORE_CODE"] = table(vlc, "codtab[TOTRUN_NUM][16] =", [15, 16], "writeSyntaxElement_Run(")
lf = JM / "app" / "lencod" / "loop_filter.h"
t["DEBLOCK_ALPHA"] = table(lf, "ALPHA_TABLE[52]", [52])
t["DEBLOCK_BETA"] = table(lf, "BETA_TABLE[52]", [52])
t["DEBLOCK_TC0"] = table(lf, "CLIP_TAB[52][5]", [52, 5])
t["CHROMA_QP"] = table(JM / "app" / "lencod" / "block.h", "QP_SCALE_CR[52]=", [52])
scan4 = table(JM / "app" / "lencod" / "block.c", "SNGL_SCAN[16][2] =", [16, 2])
t["ZIGZAG4"] = [scan4[2 * i] + 4 * scan4[2 * i + 1] for i in range(16)]
scan8 = table(JM / "app" / "lencod" / "transform8x8.c", "SNGL_SCAN8x8[64][2] =", [64, 2])
t["ZIGZAG8"] = [scan8[2 * i] + 8 * scan8[2 * i + 1] for i in range(64)]
qm = JM / "app" / "lencod" / "q_matrix.c"
t["SIG8_FRAME"] = table(FF / "h264_cabac.c", "significant_coeff_flag_offset_8x8[2][63] =", [2, 63])[:63]
cab = (FF / "cabac.c").read_text()
i = cab.index("// last_coeff_flag_offset_8x8")
t["LAST8"] = [int(x) for x in re.findall(r"\d+", cab[i + len("// last_coeff_flag_offset_8x8"):cab.index("}", i)])][:63]
ncbp = table(vlc, "NCBP[2][48][2]=", [2, 48, 2])
t["CBP_CODE_INTRA"] = [ncbp[96 + 2 * i] for i in range(48)]
t["CBP_CODE_INTER"] = [ncbp[96 + 2 * i + 1] for i in range(48)]
t["QUANT4"] = table(qm, "quant_coef[NUM_MATRIX_TYPES][4][4] =", [6, 4, 4])
t["DEQUANT4"] = table(qm, "dequant_coef[NUM_MATRIX_TYPES][4][4] =", [6, 4, 4])
t["QUANT8"] = table(qm, "quant_coef8[NUM_MATRIX_TYPES][8][8] =", [6, 8, 8])
t["DEQUANT8"] = table(qm, "dequant_coef8[NUM_MATRIX_TYPES][8][8] =", [6, 8, 8])
ctx = JM / "lib" / "lcommon" / "ctx_tables.h"
jm_mb_i = table(ctx, "INIT_MB_TYPE_I[1][3][11][2] =", [1, 3, 11, 2])
ff_i = t["CABAC_INIT_I"]
assert jm_mb_i[22:44] == ff_i[0:22], "SI/I mb_type contexts differ between JM and FFmpeg"
jm_mb_p = table(ctx, "INIT_MB_TYPE_P[3][3][11][2] =", [3, 3, 11, 2])
ff_p = t["CABAC_INIT_P"]
for idc in range(3):
row = jm_mb_p[idc * 66 + 22: idc * 66 + 44]
base = idc * 2048
assert row[0:6] == ff_p[base + 2 * 11: base + 2 * 14], f"P mb_skip_flag contexts idc {idc}"
assert row[8:22] == ff_p[base + 2 * 14: base + 2 * 21], f"P mb_type contexts idc {idc}"
assert t["ZIGZAG4"] == [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15], t["ZIGZAG4"]
assert t["QUANT4"][:3] == [13107, 8066, 13107] and t["DEQUANT4"][:3] == [10, 13, 10]
assert t["CABAC_RANGE_LPS"][:4] == [128, 176, 208, 240]
assert t["CBP_CODE_INTRA"][47] == 0 and t["CBP_CODE_INTER"][0] == 0 and t["CBP_CODE_INTRA"][0] == 3
assert len(t["LAST8"]) == 63 and t["LAST8"][-1] == 8 and t["SIG8_FRAME"][-1] == 12
assert len(t["CABAC_INIT_I"]) == 2048 and len(t["CABAC_INIT_P"]) == 6144
transition = []
for s in range(128):
state, mps = s >> 1, s & 1
for b in range(2):
if b == mps:
transition.append(t["CABAC_NEXT_MPS"][state] * 2 + mps)
else:
transition.append(t["CABAC_NEXT_LPS"][state] * 2 + (1 - mps if state == 0 else mps))
t["CABAC_TRANSITION"] = transition
sig_base = [105, 120, 134, 149, 152, 402]
last_base = [166, 181, 195, 210, 213, 417]
sig_ctx, last_ctx = [], []
for cat in range(6):
for i in range(64):
if cat == 3:
si = li = min(i, 2)
elif cat == 5:
si = t["SIG8_FRAME"][i] if i < 63 else 0
li = t["LAST8"][i] if i < 63 else 0
else:
si = li = i
sig_ctx.append(sig_base[cat] + si)
last_ctx.append(last_base[cat] + li)
t["CABAC_SIG_CTX"] = sig_ctx
t["CABAC_LAST_CTX"] = last_ctx
t["CABAC_RENORM_SHIFT"] = [6] + [next(s for s in range(9) if ((i << 3) << s) >= 256) for i in range(1, 64)]
assert t["CABAC_RENORM_SHIFT"][:4] == [6, 5, 4, 4] and t["CABAC_RENORM_SHIFT"][31] == 1
assert t["CABAC_RENORM_SHIFT"][32] == 0 and t["CABAC_TRANSITION"][0:2] == [2, 1]
parts = ["// Generated by tools/gen_tables.py from the JM 19 reference software and FFmpeg's h264_cabac.c\n"
"// (H.264 Tables 8-14..8-17, 9-5, 9-7..9-10, 9-12..9-33, 9-44). Do not edit.\n"]
for name, vals in t.items():
parts.append(emit(name, vals, elem=ELEM.get(name)))
OUT.parent.mkdir(parents=True, exist_ok=True)
OUT.write_text("\n".join(parts))
print(f"wrote {OUT} ({sum(len(v) for v in t.values())} values)")
return 0
if __name__ == "__main__":
sys.exit(main())