bef055304b
A review found rules that only the crate and the README settled, and that no vector pinned: changing the checker to another reading of each changed no committed byte. integer-operators now states each one, and each now has a vector that a different reading fails: - state: a refused decode changes no state; until a decode completes, a state input is `init`, sized by input 0 of the decode being made (§10.3). - run step 1 runs input by input; an input with a dimension of 0 is model_mismatch (§4). - run step 3: conv_transpose's uncropped output of cout planes is capped, not one plane; the caps come before the shape rule, both for that output and for an uncropped side; the shape rule comes before max_elems (crop_to) (§4, §5.3). - crop: a scalar start of −1 is refused though −1 + h ≤ H, and at load a scalar start is planned at 0, not at its range's low end (§5.10). - load step 5: chscale's m_bytes and layernorm_c's g_bytes of 2 are refused, and 4 loads (§10.4). subpos_merge now caps its output sides before its shape rule, in the crate and the checker, as §4 states for every operator. No input within the caps can tell the two orders apart. Co-Authored-By: Claude Opus 5.5 <[email protected]>
2381 lines
84 KiB
Python
2381 lines
84 KiB
Python
#!/usr/bin/env python3
|
|
"""Independent checker for the integer operator corpus (extensions/integer-operators.md).
|
|
|
|
Python standard library only; it shares no code with wai-ops. Written from the
|
|
extension's text (and, for the source forms it lowers, integer-payloads), it:
|
|
|
|
- recomputes every vector file under ops-conformance/vectors/, stage/ and graphs/:
|
|
each graph vector's `load` (WHG1 parse and load steps 1-8, section 10.5), its
|
|
`profiles` (section 14.1, from the bounds pass and verdict of sections 8.1-8.5),
|
|
and each decode's `expect` (the run of section 4, every graph operator of
|
|
section 5), and each stage vector's `expect` (section 5E); it writes each file
|
|
again canonically (sections 8.5 and 14.1) and requires the bytes of the
|
|
committed file;
|
|
- re-lowers every parameter set certificates/index.json lists, by the canonical
|
|
lowerings of section 13, and requires each lowered graph's bytes, each set's
|
|
digest (SPEC 3.1), and that every lowered graph loads; and requires the graph
|
|
of each graphs/lower_*.jsonl vector to be the role of the set its id names, as
|
|
re-lowered.
|
|
|
|
The bounds pass decides the verdict only. It does not write the certificate of
|
|
section 8.5, which the certifier half of this checker will.
|
|
|
|
python3 -I tools/wai_ops_check.py ops-conformance # check
|
|
python3 -I tools/wai_ops_check.py ops-conformance --write DIR # also write the files
|
|
|
|
Exit status 1 on any difference or failure.
|
|
"""
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
CAP = 1 << 26
|
|
I64_MIN, I64_MAX = -(1 << 63), (1 << 63) - 1
|
|
I32_MIN, I32_MAX = -(1 << 31), (1 << 31) - 1
|
|
|
|
|
|
class Refuse(Exception):
|
|
"""A refusal with a code of section 4 or 5E, and, for a stage operator, `at`."""
|
|
|
|
def __init__(self, code, at=None):
|
|
super().__init__(code)
|
|
self.code = code
|
|
self.at = at
|
|
|
|
|
|
# --------------------------------------------------------------- section 2
|
|
|
|
|
|
def R(p, s, mode):
|
|
"""R_floor, R_half_up, R_half_even (section 2)."""
|
|
if s == 0:
|
|
return p
|
|
if mode == 0:
|
|
return p >> s
|
|
if mode == 1:
|
|
return (p + (1 << (s - 1))) >> s
|
|
q = p >> s
|
|
r = p - (q << s)
|
|
h = 1 << (s - 1)
|
|
if r > h or (r == h and q % 2 == 1):
|
|
return q + 1
|
|
return q
|
|
|
|
|
|
def exp_neg_q16(x):
|
|
"""exp_neg_q16(x), x <= 0, in the order section 2 writes it."""
|
|
u = -x
|
|
n = u // 45426
|
|
g = 45426 * n
|
|
r = x + g
|
|
t = r + 88670
|
|
t2 = t * t
|
|
q = t2 >> 16
|
|
c = 23495 * q
|
|
y = (c >> 16) + 22544
|
|
return 0 if n >= 63 else y >> n
|
|
|
|
|
|
def rsqrt_q16(v):
|
|
r = math.isqrt(v)
|
|
return 0 if r == 0 else (1 << 32) // r
|
|
|
|
|
|
def trunc_div(a, b):
|
|
q = abs(a) // abs(b)
|
|
return q if (a >= 0) == (b > 0) else -q
|
|
|
|
|
|
# --------------------------------------------------------------- WHG1 (section 10)
|
|
|
|
NAMES = {
|
|
0x01: "conv", 0x02: "rescale", 0x03: "chscale", 0x04: "lut16", 0x05: "add",
|
|
0x06: "concat", 0x07: "chunk_add", 0x08: "split", 0x09: "pixel_shuffle",
|
|
0x0A: "pixel_unshuffle", 0x0B: "clamp", 0x0C: "relu", 0x0D: "leaky",
|
|
0x0E: "pad_rep", 0x0F: "crop_to", 0x10: "mul", 0x11: "conv_transpose",
|
|
0x12: "igdn", 0x13: "abs", 0x14: "crop", 0x15: "subpos_split",
|
|
0x16: "subpos_merge", 0x17: "fill", 0x18: "layernorm_c", 0x19: "l2norm",
|
|
0x1A: "softmax", 0x1B: "sigmoid", 0x1C: "elu", 0x1D: "gram", 0x1E: "apply",
|
|
}
|
|
CODE = {v: k for k, v in NAMES.items()}
|
|
TWO_IN = {"add", "crop_to", "mul", "gram", "apply"}
|
|
|
|
|
|
def arity(name):
|
|
if name in TWO_IN:
|
|
return (2, 1)
|
|
if name == "concat":
|
|
return (None, 1)
|
|
if name == "subpos_split":
|
|
return (1, 4)
|
|
if name == "subpos_merge":
|
|
return (4, 1)
|
|
return (1, 1)
|
|
|
|
|
|
class Reader:
|
|
def __init__(self, b, pos=0):
|
|
self.b, self.p = b, pos
|
|
|
|
def take(self, n):
|
|
if n < 0 or self.p + n > len(self.b):
|
|
raise Refuse("model_mismatch")
|
|
s = self.b[self.p:self.p + n]
|
|
self.p += n
|
|
return s
|
|
|
|
def u(self, n):
|
|
return int.from_bytes(self.take(n), "little")
|
|
|
|
def i(self, n):
|
|
return int.from_bytes(self.take(n), "little", signed=True)
|
|
|
|
def ints(self, count, width):
|
|
raw = self.take(count * width)
|
|
return [int.from_bytes(raw[k:k + width], "little", signed=True) for k in range(0, len(raw), width)]
|
|
|
|
def zeros(self, n):
|
|
if any(self.take(n)):
|
|
raise Refuse("model_mismatch")
|
|
|
|
|
|
def bad():
|
|
raise Refuse("model_mismatch")
|
|
|
|
|
|
def read_rq(rd, C):
|
|
"""The RQ block (section 10.4) over C channels."""
|
|
rq = {}
|
|
rq["round"], rq["bias"], rq["m_kind"], rq["s_kind"] = rd.u(1), rd.u(1), rd.u(1), rd.u(1)
|
|
rq["m_bytes"], rq["b_bytes"] = rd.u(1), rd.u(1)
|
|
if rd.u(2) != 0:
|
|
bad()
|
|
if rq["round"] > 2 or rq["bias"] > 2 or rq["m_kind"] > 2 or rq["s_kind"] not in (1, 2):
|
|
bad()
|
|
if rq["m_kind"] == 0:
|
|
if rq["m_bytes"] != 0:
|
|
bad()
|
|
elif rq["m_bytes"] not in (4, 8):
|
|
bad()
|
|
if rq["bias"] == 0:
|
|
if rq["b_bytes"] != 0:
|
|
bad()
|
|
elif rq["b_bytes"] not in (4, 8):
|
|
bad()
|
|
nm = {0: 0, 1: 1, 2: C}[rq["m_kind"]]
|
|
rq["mult"] = rd.ints(nm, rq["m_bytes"]) if nm else []
|
|
ns = 1 if rq["s_kind"] == 1 else C
|
|
rq["shift"] = list(rd.take(ns))
|
|
if any(s > 127 for s in rq["shift"]):
|
|
bad()
|
|
rq["biasv"] = rd.ints(C, rq["b_bytes"]) if rq["bias"] else []
|
|
return rq
|
|
|
|
|
|
def rq_m(rq, o):
|
|
return None if rq["m_kind"] == 0 else rq["mult"][0 if rq["m_kind"] == 1 else o]
|
|
|
|
|
|
def rq_s(rq, o):
|
|
return rq["shift"][0 if rq["s_kind"] == 1 else o]
|
|
|
|
|
|
def rq_apply(rq, o, v):
|
|
"""rq_o(v), section 5.1."""
|
|
t = v + (rq["biasv"][o] if rq["bias"] == 1 else 0)
|
|
m = rq_m(rq, o)
|
|
p = t * (1 if m is None else m)
|
|
q = R(p, rq_s(rq, o), rq["round"])
|
|
return q + (rq["biasv"][o] if rq["bias"] == 2 else 0)
|
|
|
|
|
|
def read_operand(rd, n_scalars):
|
|
kind = rd.u(1)
|
|
rd.zeros(3)
|
|
v = rd.u(4)
|
|
if kind == 0:
|
|
return ("const", v)
|
|
if kind == 1 and v < n_scalars:
|
|
return ("scalar", v)
|
|
bad()
|
|
|
|
|
|
def read_params(name, pb, cin, n_scalars):
|
|
rd = Reader(pb)
|
|
P = {}
|
|
if name == "conv":
|
|
for k in ("cout", "kh", "kw", "sh", "sw", "ph", "pw", "groups"):
|
|
P[k] = rd.u(4)
|
|
P["w_bytes"] = rd.u(1)
|
|
rd.zeros(3)
|
|
if min(P["cout"], P["kh"], P["kw"], P["sh"], P["sw"], P["groups"]) == 0 or P["w_bytes"] not in (1, 2, 4, 8):
|
|
bad()
|
|
if cin % P["groups"] or P["cout"] % P["groups"]:
|
|
bad()
|
|
P["rq"] = read_rq(rd, P["cout"])
|
|
P["w"] = rd.ints(P["cout"] * (cin // P["groups"]) * P["kh"] * P["kw"], P["w_bytes"])
|
|
elif name == "conv_transpose":
|
|
for k in ("cout", "kh", "kw", "sh", "sw", "ph", "pw", "oph", "opw"):
|
|
P[k] = rd.u(4)
|
|
P["w_bytes"] = rd.u(1)
|
|
rd.zeros(3)
|
|
if min(P["cout"], P["kh"], P["kw"], P["sh"], P["sw"]) == 0 or P["w_bytes"] not in (1, 2, 4, 8):
|
|
bad()
|
|
P["rq"] = read_rq(rd, P["cout"])
|
|
P["w"] = rd.ints(cin * P["cout"] * P["kh"] * P["kw"], P["w_bytes"])
|
|
elif name == "rescale":
|
|
P["rq"] = read_rq(rd, cin)
|
|
elif name == "chscale":
|
|
P["scalar"], P["n_rows"] = rd.u(4), rd.u(4)
|
|
P["round"], P["m_bytes"], P["shift"] = rd.u(1), rd.u(1), rd.u(1)
|
|
rd.zeros(1)
|
|
if P["scalar"] >= n_scalars or P["n_rows"] == 0 or P["m_bytes"] not in (4, 8) or P["shift"] > 127 or P["round"] > 2:
|
|
bad()
|
|
P["M"] = rd.ints(P["n_rows"] * cin, P["m_bytes"])
|
|
elif name == "mul":
|
|
P["rq"] = read_rq(rd, cin)
|
|
if P["rq"]["bias"] != 0:
|
|
bad()
|
|
elif name == "lut16":
|
|
P["e_in"], P["e_out"] = rd.i(1), rd.i(1)
|
|
rd.zeros(2)
|
|
P["T"] = rd.ints(65536, 2)
|
|
if P["T"][0] != 0:
|
|
bad()
|
|
elif name in ("add", "concat", "chunk_add", "relu", "abs", "crop_to", "softmax", "sigmoid"):
|
|
pass
|
|
elif name == "split":
|
|
P["start"], P["count"] = rd.u(4), rd.u(4)
|
|
if P["count"] == 0:
|
|
bad()
|
|
elif name in ("pixel_shuffle", "pixel_unshuffle"):
|
|
P["r"] = rd.u(4)
|
|
if P["r"] == 0:
|
|
bad()
|
|
elif name == "clamp":
|
|
P["lo"], P["hi"] = rd.i(8), rd.i(8)
|
|
if P["lo"] > P["hi"]:
|
|
bad()
|
|
elif name == "leaky":
|
|
P["a"], P["f"] = rd.i(8), rd.u(4)
|
|
if P["f"] > 63:
|
|
bad()
|
|
elif name == "pad_rep":
|
|
P["m"] = rd.u(4)
|
|
if P["m"] == 0:
|
|
bad()
|
|
elif name == "crop":
|
|
P["ops"] = [read_operand(rd, n_scalars) for _ in range(4)]
|
|
elif name in ("subpos_split", "subpos_merge"):
|
|
P["order"] = list(rd.take(4))
|
|
if sorted(P["order"]) != [0, 1, 2, 3]:
|
|
bad()
|
|
elif name == "fill":
|
|
P["kind"] = rd.u(1)
|
|
rd.zeros(7)
|
|
P["value"] = rd.i(8)
|
|
if not (P["kind"] == 0 or (P["kind"] == 1 and 0 <= P["value"] < n_scalars)):
|
|
bad()
|
|
elif name == "igdn":
|
|
P["gf"], P["frac"] = rd.u(4), rd.u(4)
|
|
P["g_bytes"] = rd.u(1)
|
|
rd.zeros(7)
|
|
if P["gf"] > 127 or P["frac"] > 127 or P["g_bytes"] not in (4, 8):
|
|
bad()
|
|
P["gamma"] = rd.ints(cin * cin, P["g_bytes"])
|
|
P["beta"] = rd.ints(cin, 8)
|
|
elif name == "layernorm_c":
|
|
P["eps"] = rd.u(8)
|
|
P["g_bytes"] = rd.u(1)
|
|
rd.zeros(7)
|
|
if P["eps"] == 0 or P["g_bytes"] not in (4, 8):
|
|
bad()
|
|
P["gamma"] = rd.ints(cin, P["g_bytes"])
|
|
P["beta"] = rd.ints(cin, 8)
|
|
elif name == "l2norm":
|
|
P["eps"] = rd.u(8)
|
|
if P["eps"] == 0:
|
|
bad()
|
|
elif name == "elu":
|
|
P["alpha"] = rd.i(8)
|
|
elif name in ("gram", "apply"):
|
|
P["groups"] = rd.u(4)
|
|
P["rq"] = read_rq(rd, cin)
|
|
if P["groups"] == 0 or P["rq"]["m_kind"] != 0 or P["rq"]["bias"] != 0 or P["rq"]["s_kind"] != 1:
|
|
bad()
|
|
if rd.p != len(pb):
|
|
bad()
|
|
return P
|
|
|
|
|
|
def parse_whg1(b):
|
|
"""Load steps 1-6 of section 10.5. Every refusal is model_mismatch."""
|
|
if len(b) < 32 or b[0:4] != b"WHG1" or b[4] != 1 or b[5] != 1:
|
|
bad()
|
|
rd = Reader(b, 6)
|
|
r0 = rd.u(2)
|
|
n_t, n_o = rd.u(4), rd.u(4)
|
|
n_i, n_out, n_s = rd.u(2), rd.u(2), rd.u(2)
|
|
r1 = rd.u(2)
|
|
max_elems = rd.u(4)
|
|
r2 = rd.u(4)
|
|
if r0 or r1 or r2:
|
|
bad()
|
|
if not (n_t and n_o and n_i and n_out) or not (1 <= max_elems <= CAP):
|
|
bad()
|
|
if 32 + 16 * n_s + 32 * n_i + 4 * n_out + 24 * n_t + 12 * n_o > len(b):
|
|
bad()
|
|
g = {"max_elems": max_elems, "scalars": [], "inputs": [], "outputs": [], "tensors": [], "ops": []}
|
|
for _ in range(n_s):
|
|
lo, hi = rd.i(8), rd.i(8)
|
|
if lo > hi:
|
|
bad()
|
|
g["scalars"].append((lo, hi))
|
|
for _ in range(n_i):
|
|
t, mh, mw = rd.u(4), rd.u(4), rd.u(4)
|
|
kind = rd.u(1)
|
|
rd.zeros(3)
|
|
src = rd.u(4)
|
|
rd.zeros(4)
|
|
init = rd.i(8)
|
|
if (mh == 0) != (mw == 0) or kind > 1 or t >= n_t:
|
|
bad()
|
|
if kind == 0 and (src or init):
|
|
bad()
|
|
g["inputs"].append({"tensor": t, "max_h": mh, "max_w": mw, "kind": kind, "src": src, "init": init})
|
|
for _ in range(n_out):
|
|
t = rd.u(4)
|
|
if t >= n_t:
|
|
bad()
|
|
g["outputs"].append(t)
|
|
for _ in range(n_t):
|
|
ch = rd.u(4)
|
|
ex = rd.i(1)
|
|
pol = rd.u(1)
|
|
rd.zeros(2)
|
|
lo, hi = rd.i(8), rd.i(8)
|
|
if ch == 0 or pol > 2:
|
|
bad()
|
|
if pol == 0 and (lo or hi):
|
|
bad()
|
|
if pol != 0 and lo > hi:
|
|
bad()
|
|
g["tensors"].append({"c": ch, "e": ex, "policy": pol, "lo": lo, "hi": hi})
|
|
for _ in range(n_o):
|
|
code, ni, no = rd.u(1), rd.u(1), rd.u(1)
|
|
rd.zeros(1)
|
|
name = NAMES.get(code)
|
|
if name is None:
|
|
bad()
|
|
ai, ao = arity(name)
|
|
if no != ao or (ai is None and ni < 2) or (ai is not None and ni != ai):
|
|
bad()
|
|
ins = [rd.u(4) for _ in range(ni)]
|
|
outs = [rd.u(4) for _ in range(no)]
|
|
if any(t >= n_t for t in ins + outs):
|
|
bad()
|
|
plen = rd.u(4)
|
|
pb = rd.take(plen)
|
|
cin = g["tensors"][ins[0]]["c"]
|
|
g["ops"].append({"name": name, "ins": ins, "outs": outs, "p": read_params(name, pb, cin, n_s)})
|
|
if rd.p != len(b):
|
|
bad()
|
|
rules(g)
|
|
return g
|
|
|
|
|
|
def rules(g):
|
|
"""Load step 6: the graph rules of section 10.3, state, exponents (3, 5.6)."""
|
|
T = g["tensors"]
|
|
defined = [False] * len(T)
|
|
for i in g["inputs"]:
|
|
if defined[i["tensor"]] or T[i["tensor"]]["policy"] == 0:
|
|
bad()
|
|
defined[i["tensor"]] = True
|
|
if g["inputs"][0]["kind"] != 0:
|
|
bad()
|
|
if len({i["max_h"] != 0 for i in g["inputs"]}) != 1:
|
|
bad()
|
|
for op in g["ops"]:
|
|
if not all(defined[t] for t in op["ins"]):
|
|
bad()
|
|
for t in op["outs"]:
|
|
if defined[t]:
|
|
bad()
|
|
defined[t] = True
|
|
if not all(defined):
|
|
bad()
|
|
for i in g["inputs"]:
|
|
if i["kind"] == 1:
|
|
if i["src"] >= len(g["outputs"]):
|
|
bad()
|
|
rec, out = T[i["tensor"]], T[g["outputs"][i["src"]]]
|
|
if out["c"] != rec["c"] or not (rec["lo"] <= i["init"] <= rec["hi"]):
|
|
bad()
|
|
for op in g["ops"]:
|
|
n, P = op["name"], op["p"]
|
|
C = T[op["ins"][0]]["c"]
|
|
cs = [T[t]["c"] for t in op["ins"]]
|
|
want = C
|
|
if n in ("conv", "conv_transpose"):
|
|
want = P["cout"]
|
|
elif n == "add":
|
|
want = C if cs[1] == C else None
|
|
elif n == "mul":
|
|
want = C if cs[1] in (C, 1) else None
|
|
elif n == "concat":
|
|
want = sum(cs)
|
|
elif n == "chunk_add":
|
|
want = C // 2 if C % 2 == 0 else None
|
|
elif n == "split":
|
|
want = P["count"] if P["start"] + P["count"] <= C else None
|
|
elif n == "pixel_shuffle":
|
|
want = C // (P["r"] ** 2) if C % (P["r"] ** 2) == 0 else None
|
|
elif n == "pixel_unshuffle":
|
|
want = C * P["r"] ** 2
|
|
elif n == "subpos_merge":
|
|
want = C if all(c == C for c in cs) else None
|
|
elif n == "fill":
|
|
want = T[op["outs"][0]]["c"]
|
|
elif n in ("gram", "apply"):
|
|
want = C if cs[1] == C and C % P["groups"] == 0 else None
|
|
elif n == "chscale":
|
|
lo, hi = g["scalars"][P["scalar"]]
|
|
want = C if lo >= 0 and hi <= P["n_rows"] - 1 else None
|
|
if want is None or any(T[t]["c"] != want for t in op["outs"]):
|
|
bad()
|
|
if n in ("add", "concat", "chunk_add"):
|
|
es = [T[t]["e"] for t in op["ins"] + op["outs"]]
|
|
if not (all(e == -128 for e in es) or (all(e != -128 for e in es) and len(set(es)) == 1)):
|
|
bad()
|
|
if n == "lut16":
|
|
if P["e_in"] != T[op["ins"][0]]["e"] or P["e_out"] != T[op["outs"][0]]["e"]:
|
|
bad()
|
|
|
|
|
|
# --------------------------------------------------------------- plans (section 4, 10.5 step 7)
|
|
|
|
|
|
class Big(Exception):
|
|
pass
|
|
|
|
|
|
class Shape(Exception):
|
|
pass
|
|
|
|
|
|
def capped(v):
|
|
if v > CAP:
|
|
raise Big()
|
|
return v
|
|
|
|
|
|
def op_sizes(g, op, hw, scal):
|
|
"""The output sizes of one operator, from its inputs' sizes. `scal` is None at load
|
|
(scalars by their ranges) or the scalar values."""
|
|
n, P = op["name"], op["p"]
|
|
h, w = hw[op["ins"][0]]
|
|
same = lambda o: hw[o] == (h, w)
|
|
if n == "conv":
|
|
hp, wp = capped(h + 2 * P["ph"]), capped(w + 2 * P["pw"])
|
|
if hp < P["kh"] or wp < P["kw"]:
|
|
raise Shape()
|
|
return [((hp - P["kh"]) // P["sh"] + 1, (wp - P["kw"]) // P["sw"] + 1)]
|
|
if n == "conv_transpose":
|
|
uh = capped((h - 1) * P["sh"] + P["kh"] + P["oph"])
|
|
uw = capped((w - 1) * P["sw"] + P["kw"] + P["opw"])
|
|
capped(P["cout"] * uh * uw)
|
|
if uh <= 2 * P["ph"] or uw <= 2 * P["pw"]:
|
|
raise Shape()
|
|
return [(uh - 2 * P["ph"], uw - 2 * P["pw"])]
|
|
if n in ("add", "mul"):
|
|
if not same(op["ins"][1]):
|
|
raise Shape()
|
|
return [(h, w)]
|
|
if n in ("concat", "subpos_merge"):
|
|
# subpos_merge's output sides are sizes it forms: capped before the shape rule.
|
|
out = (capped(2 * h), capped(2 * w)) if n == "subpos_merge" else (h, w)
|
|
if not all(same(t) for t in op["ins"]):
|
|
raise Shape()
|
|
return [out]
|
|
if n == "pixel_shuffle":
|
|
return [(capped(h * P["r"]), capped(w * P["r"]))]
|
|
if n == "pixel_unshuffle":
|
|
if h % P["r"] or w % P["r"]:
|
|
raise Shape()
|
|
return [(h // P["r"], w // P["r"])]
|
|
if n == "pad_rep":
|
|
m = P["m"]
|
|
return [(capped(-(-h // m) * m), capped(-(-w // m) * m))]
|
|
if n == "crop_to":
|
|
rh, rw = hw[op["ins"][1]]
|
|
if rh > h or rw > w:
|
|
raise Shape()
|
|
return [(rh, rw)]
|
|
if n == "crop":
|
|
y0, x0, ch, cw = P["ops"]
|
|
|
|
def side(start, length, size):
|
|
if scal is not None:
|
|
val = lambda o: o[1] if o[0] == "const" else scal[o[1]]
|
|
s0, ln = val(start), val(length)
|
|
if s0 < 0 or ln < 1 or s0 + ln > size:
|
|
raise Shape()
|
|
return ln
|
|
s0 = start[1] if start[0] == "const" else 0
|
|
ln = length[1] if length[0] == "const" else min(g["scalars"][length[1]][1], size - s0)
|
|
if ln < 1 or s0 + ln > size:
|
|
raise Shape()
|
|
return ln
|
|
|
|
return [(side(y0, ch, h), side(x0, cw, w))]
|
|
if n == "subpos_split":
|
|
if h % 2 or w % 2:
|
|
raise Shape()
|
|
return [(h // 2, w // 2)] * 4
|
|
if n == "gram":
|
|
if not same(op["ins"][1]):
|
|
raise Shape()
|
|
return [(1, g["tensors"][op["ins"][0]]["c"] // P["groups"])]
|
|
if n == "apply":
|
|
d = g["tensors"][op["ins"][0]]["c"] // P["groups"]
|
|
if (h, w) != (1, d):
|
|
raise Shape()
|
|
return [hw[op["ins"][1]]]
|
|
return [(h, w)]
|
|
|
|
|
|
def walk(g, hw, scal, load):
|
|
for op in g["ops"]:
|
|
try:
|
|
outs = op_sizes(g, op, hw, scal)
|
|
except Big:
|
|
raise Refuse("too_large")
|
|
except Shape:
|
|
raise Refuse("model_mismatch" if load else "geometry_mismatch")
|
|
for t, (oh, ow) in zip(op["outs"], outs):
|
|
if g["tensors"][t]["c"] * oh * ow > g["max_elems"]:
|
|
raise Refuse("too_large")
|
|
hw[t] = (oh, ow)
|
|
|
|
|
|
def load(b):
|
|
"""Load steps 1-8 of section 10.5. Step 8 is the bounds pass and the verdict
|
|
(section 8): a graph whose verdict meets neither profile is refused, so a graph
|
|
that loads meets acc64 at least, and g["verdict"] says whether it meets acc32."""
|
|
g = parse_whg1(b)
|
|
if g["inputs"][0]["max_h"] == 0:
|
|
if any(op["name"] in ("gram", "l2norm", "softmax") for op in g["ops"]):
|
|
bad()
|
|
g["sizes"] = None
|
|
else:
|
|
hw = {}
|
|
for i in g["inputs"]:
|
|
c = g["tensors"][i["tensor"]]["c"]
|
|
if c * i["max_h"] * i["max_w"] > g["max_elems"]:
|
|
raise Refuse("too_large")
|
|
hw[i["tensor"]] = (i["max_h"], i["max_w"])
|
|
walk(g, hw, None, True)
|
|
g["sizes"] = hw
|
|
g["verdict"] = bounds(g)
|
|
if not g["verdict"]["acc64"]:
|
|
bad()
|
|
return g
|
|
|
|
|
|
# --------------------------------------------------------------- bounds (section 8)
|
|
|
|
SAT = 1 << 255
|
|
# Section 7.1: the width of each class under each profile, in bits.
|
|
WIDTH = {
|
|
"acc64": {"tensor": 64, "acc": 64, "wide": 128, "product": 128},
|
|
"acc32": {"tensor": 32, "acc": 32, "wide": 32, "product": 64},
|
|
}
|
|
# The class of each named intermediate (sections 5 and 2).
|
|
CLASS = {
|
|
"conv.term": "acc", "conv.acc": "acc", "convT.term": "acc", "convT.acc": "acc",
|
|
"rq.t": "acc", "rq.p": "product", "rq.q": "acc", "rq.y": "acc",
|
|
"mul.p": "acc", "leaky.p": "acc",
|
|
"gram.term": "acc", "gram.acc": "acc", "apply.term": "acc", "apply.acc": "acc",
|
|
"igdn.x2": "wide", "igdn.gx2": "product", "igdn.term": "wide", "igdn.N": "wide",
|
|
"igdn.r": "acc", "igdn.xr": "wide",
|
|
"ln.S": "wide", "ln.T": "wide", "ln.m": "wide", "ln.v": "wide", "ln.dk": "wide", "ln.ng": "wide",
|
|
"ln.d": "acc", "rsq.num": "acc", "ln.k": "acc", "ln.n": "acc",
|
|
"l2.T": "wide", "l2.vk": "wide", "l2.k": "acc",
|
|
"sm.diff": "acc", "sm.sum": "acc", "sm.num": "acc", "sig.num": "acc", "elu.p": "wide",
|
|
"exp.negx": "acc", "exp.nl": "acc", "exp.t2": "acc", "exp.c": "acc",
|
|
}
|
|
|
|
|
|
def fits_bits(iv, n):
|
|
return -(1 << (n - 1)) <= iv[0] and iv[1] <= (1 << (n - 1)) - 1
|
|
|
|
|
|
def hull(ivs):
|
|
return (min(v[0] for v in ivs), max(v[1] for v in ivs))
|
|
|
|
|
|
def isqrt_ceil(n):
|
|
r = math.isqrt(n)
|
|
return r if r * r == n else r + 1
|
|
|
|
|
|
class Node:
|
|
"""One node of the bounds pass. Every +, -, x and floor division (a shift
|
|
included) goes through sat (section 8.1), and the node is saturated when sat
|
|
changes any result. min, max, |.|, negation and isqrt cannot leave
|
|
[-2^255, 2^255] when their operands lie in it. `rec` holds each named
|
|
intermediate's interval, hulled over channels."""
|
|
|
|
def __init__(self):
|
|
self.saturated = False
|
|
self.rec = {}
|
|
|
|
def sat(self, v):
|
|
if v > SAT:
|
|
self.saturated = True
|
|
return SAT
|
|
if v < -SAT:
|
|
self.saturated = True
|
|
return -SAT
|
|
return v
|
|
|
|
def add(self, a, b):
|
|
return self.sat(a + b)
|
|
|
|
def sub(self, a, b):
|
|
return self.sat(a - b)
|
|
|
|
def mul(self, a, b):
|
|
return self.sat(a * b)
|
|
|
|
def div(self, a, b):
|
|
return self.sat(a // b)
|
|
|
|
def shr(self, a, s):
|
|
return self.sat(a >> s)
|
|
|
|
def note(self, key, iv):
|
|
old = self.rec.get(key)
|
|
self.rec[key] = iv if old is None else (min(old[0], iv[0]), max(old[1], iv[1]))
|
|
return iv
|
|
|
|
def term(self, w, x):
|
|
"""The term interval of a weight w over x (section 8.2)."""
|
|
a, b = self.mul(w, x[0]), self.mul(w, x[1])
|
|
return (min(a, b), max(a, b))
|
|
|
|
def corner(self, x, y):
|
|
c = [self.mul(x[0], y[0]), self.mul(x[0], y[1]), self.mul(x[1], y[0]), self.mul(x[1], y[1])]
|
|
return (min(c), max(c))
|
|
|
|
def A(self, terms, b):
|
|
"""The partial-sum interval A(K, b) of section 8.2."""
|
|
lo, hi = min(0, b), max(0, b)
|
|
for t in terms:
|
|
lo = self.add(lo, min(0, t[0]))
|
|
hi = self.add(hi, max(0, t[1]))
|
|
return (lo, hi)
|
|
|
|
def R(self, p, s, mode):
|
|
"""R_floor, R_half_up, R_half_even (section 2), as written there."""
|
|
if s == 0:
|
|
return p
|
|
if mode == 0:
|
|
return self.shr(p, s)
|
|
if mode == 1:
|
|
return self.shr(self.add(p, 1 << (s - 1)), s)
|
|
q = self.shr(p, s)
|
|
r = self.sub(p, self.mul(q, 1 << s))
|
|
h = 1 << (s - 1)
|
|
return self.add(q, 1) if r > h or (r == h and q % 2 == 1) else q
|
|
|
|
def rq(self, t, m, s, mode, post):
|
|
"""The rq.p, rq.q and rq.y rules of section 8.3 for one channel."""
|
|
self.note("rq.t", t)
|
|
P = self.term(m, t)
|
|
a = (1 << (s - 1)) if mode != 0 and s >= 1 else 0
|
|
self.note("rq.p", (P[0], self.add(P[1], a)))
|
|
q = (self.R(P[0], s, mode), self.R(P[1], s, mode))
|
|
self.note("rq.q", q)
|
|
return self.note("rq.y", (self.add(q[0], post), self.add(q[1], post)))
|
|
|
|
def exp(self, u):
|
|
"""The exp_neg_q16 rule of section 8.3, given the interval of u = -x."""
|
|
self.note("exp.negx", u)
|
|
self.note("exp.nl", (self.mul(45426, self.div(u[0], 45426)), self.mul(45426, self.div(u[1], 45426))))
|
|
self.note("exp.t2", (self.mul(43245, 43245), self.mul(88670, 88670)))
|
|
self.note("exp.c", (0, self.mul(23495, self.div(self.mul(88670, 88670), 1 << 16))))
|
|
|
|
|
|
def bpre(rq, o):
|
|
return rq["biasv"][o] if rq["bias"] == 1 else 0
|
|
|
|
|
|
def bpost(rq, o):
|
|
return rq["biasv"][o] if rq["bias"] == 2 else 0
|
|
|
|
|
|
def rq_mult(rq, o):
|
|
m = rq_m(rq, o)
|
|
return 1 if m is None else m
|
|
|
|
|
|
def node_bounds(nb, g, op, iv):
|
|
"""The producer results of one operator (section 8.3), recording its named
|
|
intermediates in nb. Refuses (model_mismatch) when lut16's input range is not
|
|
within [-32767, 32767] (section 5.6)."""
|
|
T = g["tensors"]
|
|
n, P = op["name"], op["p"]
|
|
x = iv[op["ins"][0]]
|
|
l, h = x
|
|
C = T[op["ins"][0]]["c"]
|
|
if n == "conv":
|
|
rq, cg, kk = P["rq"], C // P["groups"], P["kh"] * P["kw"]
|
|
outs = []
|
|
for o in range(P["cout"]):
|
|
taps = [nb.term(P["w"][(o * cg + ci) * kk + k], x) for ci in range(cg) for k in range(kk)]
|
|
nb.note("conv.term", hull(taps))
|
|
acc = nb.note("conv.acc", nb.A(taps, bpre(rq, o)))
|
|
outs.append(nb.rq(acc, rq_mult(rq, o), rq_s(rq, o), rq["round"], bpost(rq, o)))
|
|
return [hull(outs)]
|
|
if n == "conv_transpose":
|
|
rq, cout, kh, kw, sh, sw = P["rq"], P["cout"], P["kh"], P["kw"], P["sh"], P["sw"]
|
|
outs = []
|
|
for o in range(cout):
|
|
every, by_res = [], {}
|
|
for ci in range(C):
|
|
for i in range(kh):
|
|
for j in range(kw):
|
|
t = nb.term(P["w"][((ci * cout + o) * kh + i) * kw + j], x)
|
|
every.append(t)
|
|
by_res.setdefault((i % sh, j % sw), []).append(t)
|
|
nb.note("convT.term", hull(every))
|
|
b = bpre(rq, o)
|
|
accs = [nb.A(by_res[r], b) for r in sorted(by_res)]
|
|
if len(by_res) < sh * sw:
|
|
# A residue that no tap has: A over no terms.
|
|
accs.append(nb.A([], b))
|
|
acc = nb.note("convT.acc", hull(accs))
|
|
outs.append(nb.rq(acc, rq_mult(rq, o), rq_s(rq, o), rq["round"], bpost(rq, o)))
|
|
return [hull(outs)]
|
|
if n == "rescale":
|
|
rq = P["rq"]
|
|
t = lambda c: (nb.add(l, bpre(rq, c)), nb.add(h, bpre(rq, c)))
|
|
return [hull([nb.rq(t(c), rq_mult(rq, c), rq_s(rq, c), rq["round"], bpost(rq, c)) for c in range(C)])]
|
|
if n == "chscale":
|
|
lo_r, hi_r = g["scalars"][P["scalar"]]
|
|
s, mode = P["shift"], P["round"]
|
|
a = (1 << (s - 1)) if mode != 0 and s >= 1 else 0
|
|
outs = []
|
|
for c in range(C):
|
|
ps, qs = [], []
|
|
for r in range(lo_r, hi_r + 1):
|
|
Pm = nb.term(P["M"][r * C + c], x)
|
|
ps.append((Pm[0], nb.add(Pm[1], a)))
|
|
qs.append((nb.R(Pm[0], s, mode), nb.R(Pm[1], s, mode)))
|
|
nb.note("rq.t", x)
|
|
nb.note("rq.p", hull(ps))
|
|
q = nb.note("rq.q", hull(qs))
|
|
outs.append(nb.note("rq.y", q))
|
|
return [hull(outs)]
|
|
if n == "mul":
|
|
rq = P["rq"]
|
|
mp = nb.note("mul.p", nb.corner(x, iv[op["ins"][1]]))
|
|
return [hull([nb.rq(mp, rq_mult(rq, c), rq_s(rq, c), rq["round"], 0) for c in range(C)])]
|
|
if n in ("gram", "apply"):
|
|
rq = P["rq"]
|
|
t = nb.note(n + ".term", nb.corner(x, iv[op["ins"][1]]))
|
|
if n == "gram":
|
|
L = max(g["sizes"][u][0] * g["sizes"][u][1] for u in op["ins"])
|
|
else:
|
|
L = C // P["groups"]
|
|
acc = nb.note(n + ".acc", (nb.mul(L, min(0, t[0])), nb.mul(L, max(0, t[1]))))
|
|
return [nb.rq(acc, 1, rq_s(rq, 0), rq["round"], 0)]
|
|
if n == "add":
|
|
b = iv[op["ins"][1]]
|
|
return [(nb.add(l, b[0]), nb.add(h, b[1]))]
|
|
if n == "chunk_add":
|
|
return [(nb.mul(2, l), nb.mul(2, h))]
|
|
if n == "clamp":
|
|
return [(min(max(l, P["lo"]), P["hi"]), min(max(h, P["lo"]), P["hi"]))]
|
|
if n == "relu":
|
|
return [(max(l, 0), max(h, 0))]
|
|
if n == "abs":
|
|
return [(l, h) if l >= 0 else ((-h, -l) if h <= 0 else (0, max(-l, h)))]
|
|
if n == "leaky":
|
|
parts = []
|
|
if l < 0:
|
|
p = nb.note("leaky.p", nb.corner((l, min(h, -1)), (P["a"], P["a"])))
|
|
parts.append((nb.shr(p[0], P["f"]), nb.shr(p[1], P["f"])))
|
|
else:
|
|
nb.note("leaky.p", (0, 0))
|
|
if h >= 0:
|
|
parts.append((max(l, 0), h))
|
|
return [hull(parts)]
|
|
if n == "lut16":
|
|
if l < -32767 or h > 32767:
|
|
bad()
|
|
vals = P["T"][l + 32768:h + 32769]
|
|
return [(min(vals), max(vals))]
|
|
if n == "igdn":
|
|
l2, h2 = nb.mul(l, l), nb.mul(h, h)
|
|
x2 = nb.note("igdn.x2", (0, max(l2, h2)) if l <= 0 <= h else (min(l2, h2), max(l2, h2)))
|
|
outs = []
|
|
for i in range(C):
|
|
terms = []
|
|
for j in range(C):
|
|
gm = P["gamma"][i * C + j]
|
|
gx = nb.note("igdn.gx2", nb.corner((gm, gm), x2))
|
|
terms.append(nb.note("igdn.term", (nb.shr(gx[0], P["gf"]), nb.shr(gx[1], P["gf"]))))
|
|
nb.note("igdn.N", nb.A(terms, P["beta"][i]))
|
|
# The final norm, a left fold over j from beta_i; not recorded.
|
|
f_lo = f_hi = P["beta"][i]
|
|
for t in terms:
|
|
f_lo, f_hi = nb.add(f_lo, t[0]), nb.add(f_hi, t[1])
|
|
r = nb.note("igdn.r", (math.isqrt(max(0, f_lo)), math.isqrt(max(0, f_hi))))
|
|
xr = nb.note("igdn.xr", nb.corner(x, r))
|
|
outs.append((nb.shr(xr[0], P["frac"]), nb.shr(xr[1], P["frac"])))
|
|
return [hull(outs)]
|
|
if n == "layernorm_c":
|
|
e = P["eps"]
|
|
r0 = math.isqrt(e)
|
|
nb.note("ln.S", (nb.mul(C, min(0, l)), nb.mul(C, max(0, h))))
|
|
half = C // 2
|
|
nb.note("ln.m", (nb.sub(nb.mul(C, l), half), nb.add(nb.mul(C, h), half)))
|
|
d = nb.sub(h, l)
|
|
nb.note("ln.d", (nb.sub(l, h), d))
|
|
dd = nb.mul(d, d)
|
|
nb.note("ln.T", (0, nb.mul(C, dd)))
|
|
nb.note("ln.v", (e, nb.add(dd, e)))
|
|
nb.note("rsq.num", (1 << 32, 1 << 32))
|
|
k = nb.div(1 << 32, r0)
|
|
nb.note("ln.k", (0, k))
|
|
ic = isqrt_ceil(C)
|
|
big = nb.mul(nb.mul(ic, 1 << 32), nb.add(r0, 1))
|
|
D = min(nb.mul(d, k), -nb.div(-big, r0))
|
|
nb.note("ln.dk", (-D, D))
|
|
nmax = nb.add(nb.div(nb.mul(nb.mul(ic, 1 << 16), nb.add(r0, 1)), r0), 1)
|
|
N = min(nb.add(nb.shr(nb.mul(d, k), 16), 1), nmax)
|
|
nb.note("ln.n", (-N, N))
|
|
outs = []
|
|
for c in range(C):
|
|
gm = P["gamma"][c]
|
|
ng = nb.note("ln.ng", nb.corner((-N, N), (gm, gm)))
|
|
outs.append((nb.add(nb.shr(ng[0], 16), P["beta"][c]), nb.add(nb.shr(ng[1], 16), P["beta"][c])))
|
|
return [hull(outs)]
|
|
if n == "l2norm":
|
|
e = P["eps"]
|
|
hh, ww = g["sizes"][op["ins"][0]]
|
|
X = max(abs(l), abs(h))
|
|
nb.note("l2.T", (0, nb.add(nb.mul(nb.mul(hh, ww), nb.mul(X, X)), e)))
|
|
nb.note("rsq.num", (1 << 32, 1 << 32))
|
|
nb.note("l2.k", (0, nb.div(1 << 32, math.isqrt(e))))
|
|
nb.note("l2.vk", (-(1 << 32), 1 << 32))
|
|
return [(-(1 << 16), 1 << 16)]
|
|
if n == "softmax":
|
|
nb.note("sm.diff", (nb.sub(l, h), 0))
|
|
nb.exp((0, nb.sub(h, l)))
|
|
wm = g["sizes"][op["ins"][0]][1]
|
|
nb.note("sm.sum", (0, nb.mul(wm, 65553)))
|
|
nb.note("sm.num", (0, nb.add(nb.mul(65553, 1 << 16), nb.div(nb.mul(wm, 65553), 2))))
|
|
return [(0, 1 << 16)]
|
|
if n == "sigmoid":
|
|
nb.exp((l, h) if l >= 0 else ((-h, -l) if h <= 0 else (0, max(-l, h))))
|
|
nb.note("sig.num", (0, nb.mul(65553, 1 << 16)))
|
|
return [(0, 1 << 16)]
|
|
if n == "elu":
|
|
nb.exp((max(0, -h), max(0, -l)))
|
|
al = P["alpha"]
|
|
p = nb.note("elu.p", nb.corner((al, al), (-(1 << 16), 17)))
|
|
parts = []
|
|
if h > 0:
|
|
parts.append((max(l, 1), h))
|
|
if l <= 0:
|
|
parts.append((nb.shr(p[0], 16), nb.shr(p[1], 16)))
|
|
return [hull(parts)]
|
|
if n in ("concat", "subpos_merge"):
|
|
return [hull([iv[t] for t in op["ins"]])]
|
|
if n == "subpos_split":
|
|
return [x] * 4
|
|
if n == "fill":
|
|
return [(P["value"], P["value"]) if P["kind"] == 0 else g["scalars"][P["value"]]]
|
|
# split, pixel_shuffle, pixel_unshuffle, pad_rep, crop_to (its first input), crop.
|
|
return [x]
|
|
|
|
|
|
def params32(op):
|
|
"""The parameter conditions of section 7.2: whether the node meets every one."""
|
|
n, P = op["name"], op["p"]
|
|
f32 = lambda v: I32_MIN <= v <= I32_MAX
|
|
ok = True
|
|
rq = P.get("rq")
|
|
if rq is not None:
|
|
ok = all(map(f32, rq["mult"])) and all(s <= 63 for s in rq["shift"]) and (rq["bias"] != 2 or all(map(f32, rq["biasv"])))
|
|
if n in ("conv", "conv_transpose"):
|
|
ok = ok and all(map(f32, P["w"]))
|
|
elif n == "chscale":
|
|
ok = all(map(f32, P["M"])) and P["shift"] <= 63
|
|
elif n == "igdn":
|
|
ok = all(map(f32, P["gamma"])) and P["gf"] <= 63 and P["frac"] <= 63
|
|
elif n == "leaky":
|
|
ok = f32(P["a"])
|
|
return ok
|
|
|
|
|
|
def bounds(g):
|
|
"""The bounds pass (sections 8.1-8.4) and the verdict (section 8.5) of a graph
|
|
that passed load steps 1-7: {"acc32": bool, "acc64": bool}. Refuses
|
|
(model_mismatch) where the pass itself refuses the graph at load: a `check`
|
|
interval that is empty (section 8.1), or lut16's input range (section 5.6)."""
|
|
T = g["tensors"]
|
|
ok = {"acc32": True, "acc64": True}
|
|
iv = {}
|
|
declared = [(T[i["tensor"]]["lo"], T[i["tensor"]]["hi"]) for i in g["inputs"]] + list(g["scalars"])
|
|
for r in declared:
|
|
for p in ok:
|
|
ok[p] = ok[p] and fits_bits(r, WIDTH[p]["tensor"])
|
|
for i in g["inputs"]:
|
|
iv[i["tensor"]] = (T[i["tensor"]]["lo"], T[i["tensor"]]["hi"])
|
|
for op in g["ops"]:
|
|
nb = Node()
|
|
results = node_bounds(nb, g, op, iv)
|
|
fit = {p: not nb.saturated for p in ok}
|
|
for p in ok:
|
|
for k, r in nb.rec.items():
|
|
fit[p] = fit[p] and fits_bits(r, WIDTH[p][CLASS[k]])
|
|
for r in results:
|
|
fit[p] = fit[p] and fits_bits(r, WIDTH[p]["tensor"])
|
|
fit["acc32"] = fit["acc32"] and params32(op)
|
|
for t, (a, b) in zip(op["outs"], results):
|
|
rec = T[t]
|
|
if rec["policy"] == 0:
|
|
iv[t] = (a, b)
|
|
elif rec["policy"] == 1:
|
|
iv[t] = (min(max(a, rec["lo"]), rec["hi"]), min(max(b, rec["lo"]), rec["hi"]))
|
|
else:
|
|
fit["acc32"] = False
|
|
iv[t] = (max(a, rec["lo"]), min(b, rec["hi"]))
|
|
if iv[t][0] > iv[t][1]:
|
|
bad()
|
|
for p in ok:
|
|
ok[p] = ok[p] and fit[p]
|
|
return ok
|
|
|
|
|
|
# --------------------------------------------------------------- the operators (section 5)
|
|
|
|
|
|
class T:
|
|
"""A tensor [c, h, w], elements in CHW order."""
|
|
|
|
def __init__(self, c, h, w, d):
|
|
self.c, self.h, self.w, self.d = c, h, w, d
|
|
|
|
def at(self, c, y, x):
|
|
return self.d[(c * self.h + y) * self.w + x]
|
|
|
|
|
|
def ev_conv(x, P, oh, ow):
|
|
cin, H, W = x.c, x.h, x.w
|
|
G = P["groups"]
|
|
cg, og = cin // G, P["cout"] // G
|
|
kh, kw = P["kh"], P["kw"]
|
|
d = []
|
|
for o in range(P["cout"]):
|
|
base = (o // og) * cg
|
|
for y in range(oh):
|
|
for xx in range(ow):
|
|
acc = 0
|
|
for ci in range(cg):
|
|
for i in range(kh):
|
|
u = y * P["sh"] + i - P["ph"]
|
|
if not 0 <= u < H:
|
|
continue
|
|
for j in range(kw):
|
|
v = xx * P["sw"] + j - P["pw"]
|
|
if 0 <= v < W:
|
|
acc += P["w"][((o * cg + ci) * kh + i) * kw + j] * x.at(base + ci, u, v)
|
|
d.append(rq_apply(P["rq"], o, acc))
|
|
return T(P["cout"], oh, ow, d)
|
|
|
|
|
|
def ev_conv_t(x, P, oh, ow):
|
|
"""The gather form of section 5.3."""
|
|
cin, H, W = x.c, x.h, x.w
|
|
kh, kw, sh, sw, cout = P["kh"], P["kw"], P["sh"], P["sw"], P["cout"]
|
|
d = []
|
|
for o in range(cout):
|
|
for y in range(oh):
|
|
for xx in range(ow):
|
|
acc = 0
|
|
for i in range(kh):
|
|
ty = y + P["ph"] - i
|
|
if ty < 0 or ty % sh or ty // sh >= H:
|
|
continue
|
|
for j in range(kw):
|
|
tx = xx + P["pw"] - j
|
|
if tx < 0 or tx % sw or tx // sw >= W:
|
|
continue
|
|
for ci in range(cin):
|
|
acc += P["w"][((ci * cout + o) * kh + i) * kw + j] * x.at(ci, ty // sh, tx // sw)
|
|
d.append(rq_apply(P["rq"], o, acc))
|
|
return T(cout, oh, ow, d)
|
|
|
|
|
|
def per_channel(x, f):
|
|
plane = x.h * x.w
|
|
return T(x.c, x.h, x.w, [f(k // plane, v) for k, v in enumerate(x.d)])
|
|
|
|
|
|
def ev(g, op, vals, hw, scal):
|
|
n, P = op["name"], op["p"]
|
|
x = vals[op["ins"][0]]
|
|
oh, ow = hw[op["outs"][0]]
|
|
oc = [g["tensors"][t]["c"] for t in op["outs"]]
|
|
if n == "conv":
|
|
return [ev_conv(x, P, oh, ow)]
|
|
if n == "conv_transpose":
|
|
return [ev_conv_t(x, P, oh, ow)]
|
|
if n == "rescale":
|
|
return [per_channel(x, lambda c, v: rq_apply(P["rq"], c, v))]
|
|
if n == "chscale":
|
|
r = scal[P["scalar"]]
|
|
return [per_channel(x, lambda c, v: R(v * P["M"][r * x.c + c], P["shift"], P["round"]))]
|
|
if n == "mul":
|
|
b = vals[op["ins"][1]]
|
|
plane = x.h * x.w
|
|
d = []
|
|
for c in range(x.c):
|
|
cb = c if b.c == x.c else 0
|
|
for k in range(plane):
|
|
d.append(rq_apply(P["rq"], c, x.d[c * plane + k] * b.d[cb * plane + k]))
|
|
return [T(x.c, x.h, x.w, d)]
|
|
if n == "lut16":
|
|
return [T(x.c, x.h, x.w, [P["T"][v + 32768] if -32768 <= v <= 32767 else None for v in x.d])]
|
|
if n == "add":
|
|
b = vals[op["ins"][1]]
|
|
return [T(x.c, x.h, x.w, [p + q for p, q in zip(x.d, b.d)])]
|
|
if n == "concat":
|
|
d = []
|
|
for t in op["ins"]:
|
|
d += vals[t].d
|
|
return [T(oc[0], x.h, x.w, d)]
|
|
if n == "chunk_add":
|
|
half = (x.c // 2) * x.h * x.w
|
|
return [T(x.c // 2, x.h, x.w, [x.d[k] + x.d[k + half] for k in range(half)])]
|
|
if n == "split":
|
|
plane = x.h * x.w
|
|
return [T(P["count"], x.h, x.w, x.d[P["start"] * plane:(P["start"] + P["count"]) * plane])]
|
|
if n == "pixel_shuffle":
|
|
r = P["r"]
|
|
co = x.c // (r * r)
|
|
d = [0] * (co * oh * ow)
|
|
for c in range(co):
|
|
for i in range(r):
|
|
for j in range(r):
|
|
for hh in range(x.h):
|
|
for ww in range(x.w):
|
|
d[(c * oh + hh * r + i) * ow + ww * r + j] = x.at(c * r * r + i * r + j, hh, ww)
|
|
return [T(co, oh, ow, d)]
|
|
if n == "pixel_unshuffle":
|
|
r = P["r"]
|
|
d = [0] * (x.c * r * r * oh * ow)
|
|
for c in range(x.c):
|
|
for i in range(r):
|
|
for j in range(r):
|
|
for hh in range(oh):
|
|
for ww in range(ow):
|
|
d[((c * r * r + i * r + j) * oh + hh) * ow + ww] = x.at(c, hh * r + i, ww * r + j)
|
|
return [T(x.c * r * r, oh, ow, d)]
|
|
if n == "clamp":
|
|
return [T(x.c, x.h, x.w, [min(max(v, P["lo"]), P["hi"]) for v in x.d])]
|
|
if n == "relu":
|
|
return [T(x.c, x.h, x.w, [max(v, 0) for v in x.d])]
|
|
if n == "leaky":
|
|
return [T(x.c, x.h, x.w, [v if v >= 0 else (v * P["a"]) >> P["f"] for v in x.d])]
|
|
if n == "abs":
|
|
return [T(x.c, x.h, x.w, [abs(v) for v in x.d])]
|
|
if n == "pad_rep":
|
|
return [T(x.c, oh, ow, [x.at(c, min(y, x.h - 1), min(xx, x.w - 1)) for c in range(x.c) for y in range(oh) for xx in range(ow)])]
|
|
if n in ("crop_to", "crop"):
|
|
y0 = x0 = 0
|
|
if n == "crop":
|
|
val = lambda o: o[1] if o[0] == "const" else scal[o[1]]
|
|
y0, x0 = val(P["ops"][0]), val(P["ops"][1])
|
|
return [T(x.c, oh, ow, [x.at(c, y0 + y, x0 + xx) for c in range(x.c) for y in range(oh) for xx in range(ow)])]
|
|
if n == "subpos_split":
|
|
outs = []
|
|
for k in range(4):
|
|
a, b2 = P["order"][k] >> 1, P["order"][k] & 1
|
|
outs.append(T(x.c, oh, ow, [x.at(c, 2 * y + a, 2 * xx + b2) for c in range(x.c) for y in range(oh) for xx in range(ow)]))
|
|
return outs
|
|
if n == "subpos_merge":
|
|
d = [0] * (x.c * oh * ow)
|
|
for k in range(4):
|
|
a, b2 = P["order"][k] >> 1, P["order"][k] & 1
|
|
t = vals[op["ins"][k]]
|
|
for c in range(x.c):
|
|
for y in range(t.h):
|
|
for xx in range(t.w):
|
|
d[(c * oh + 2 * y + a) * ow + 2 * xx + b2] = t.at(c, y, xx)
|
|
return [T(x.c, oh, ow, d)]
|
|
if n == "fill":
|
|
v = P["value"] if P["kind"] == 0 else scal[P["value"]]
|
|
return [T(oc[0], x.h, x.w, [v] * (oc[0] * x.h * x.w))]
|
|
if n == "igdn":
|
|
C, plane = x.c, x.h * x.w
|
|
out = [0] * (C * plane)
|
|
for p in range(plane):
|
|
xs = [x.d[j * plane + p] for j in range(C)]
|
|
for i in range(C):
|
|
N = P["beta"][i] + sum((P["gamma"][i * C + j] * xs[j] * xs[j]) // (1 << P["gf"]) for j in range(C))
|
|
r = math.isqrt(max(0, N))
|
|
out[i * plane + p] = (xs[i] * r) // (1 << P["frac"])
|
|
return [T(C, x.h, x.w, out)]
|
|
if n == "layernorm_c":
|
|
C, plane = x.c, x.h * x.w
|
|
out = [0] * (C * plane)
|
|
for p in range(plane):
|
|
xs = [x.d[c * plane + p] for c in range(C)]
|
|
S = sum(xs)
|
|
sg = (S > 0) - (S < 0)
|
|
M = S + (C // 2) * sg
|
|
m = trunc_div(M, C)
|
|
ds = [v - m for v in xs]
|
|
V = sum(d * d for d in ds) // C
|
|
k = rsqrt_q16(V + P["eps"])
|
|
for c in range(C):
|
|
nn = (ds[c] * k) >> 16
|
|
out[c * plane + p] = ((nn * P["gamma"][c]) >> 16) + P["beta"][c]
|
|
return [T(C, x.h, x.w, out)]
|
|
if n == "l2norm":
|
|
plane = x.h * x.w
|
|
out = []
|
|
for c in range(x.c):
|
|
vs = x.d[c * plane:(c + 1) * plane]
|
|
k = rsqrt_q16(sum(v * v for v in vs) + P["eps"])
|
|
out += [(v * k) >> 16 for v in vs]
|
|
return [T(x.c, x.h, x.w, out)]
|
|
if n == "softmax":
|
|
out = []
|
|
for row in range(x.c * x.h):
|
|
vs = x.d[row * x.w:(row + 1) * x.w]
|
|
m = max(vs)
|
|
es = [exp_neg_q16(v - m) for v in vs]
|
|
s = sum(es)
|
|
out += [(e * 65536 + s // 2) // s for e in es]
|
|
return [T(x.c, x.h, x.w, out)]
|
|
if n == "sigmoid":
|
|
def sig(v):
|
|
if v >= 0:
|
|
return (1 << 32) // (65536 + exp_neg_q16(-v))
|
|
e = exp_neg_q16(v)
|
|
return (e * 65536) // (65536 + e)
|
|
return [T(x.c, x.h, x.w, [sig(v) for v in x.d])]
|
|
if n == "elu":
|
|
return [T(x.c, x.h, x.w, [v if v > 0 else (P["alpha"] * (exp_neg_q16(v) - 65536)) >> 16 for v in x.d])]
|
|
if n == "gram":
|
|
b = vals[op["ins"][1]]
|
|
G = P["groups"]
|
|
dd = x.c // G
|
|
plane = x.h * x.w
|
|
out = [0] * (x.c * dd)
|
|
for gi in range(G):
|
|
for i in range(dd):
|
|
for j in range(dd):
|
|
s = sum(x.d[(gi * dd + i) * plane + p] * b.d[(gi * dd + j) * plane + p] for p in range(plane))
|
|
out[(gi * dd + i) * dd + j] = R(s, P["rq"]["shift"][0], P["rq"]["round"])
|
|
return [T(x.c, 1, dd, out)]
|
|
if n == "apply":
|
|
V = vals[op["ins"][1]]
|
|
G = P["groups"]
|
|
dd = x.c // G
|
|
plane = V.h * V.w
|
|
out = [0] * (V.c * plane)
|
|
for gi in range(G):
|
|
for i in range(dd):
|
|
for p in range(plane):
|
|
s = sum(x.d[(gi * dd + i) * dd + j] * V.d[(gi * dd + j) * plane + p] for j in range(dd))
|
|
out[(gi * dd + i) * plane + p] = R(s, P["rq"]["shift"][0], P["rq"]["round"])
|
|
return [T(V.c, V.h, V.w, out)]
|
|
raise Refuse("model_mismatch")
|
|
|
|
|
|
def apply_policy(rec, t):
|
|
"""Returns False when a `check` tensor holds an element outside its range."""
|
|
if rec["policy"] == 1:
|
|
t.d = [min(max(v, rec["lo"]), rec["hi"]) for v in t.d]
|
|
elif rec["policy"] == 2:
|
|
return all(rec["lo"] <= v <= rec["hi"] for v in t.d)
|
|
return True
|
|
|
|
|
|
def run(g, data_inputs, scalars, state):
|
|
"""One decode (section 4): the data inputs in input-table order; each state input
|
|
from `state` (section 10.3), which the decode updates when it completes."""
|
|
allin = []
|
|
di = 0
|
|
si = 0
|
|
if len(data_inputs) != sum(1 for i in g["inputs"] if i["kind"] == 0):
|
|
raise Refuse("model_mismatch")
|
|
for i in g["inputs"]:
|
|
if i["kind"] == 0:
|
|
allin.append(data_inputs[di])
|
|
di += 1
|
|
else:
|
|
if state.get("s") is None:
|
|
x0 = data_inputs[0]
|
|
c = g["tensors"][i["tensor"]]["c"]
|
|
allin.append(T(c, x0.h, x0.w, [i["init"]] * (c * x0.h * x0.w)))
|
|
else:
|
|
allin.append(state["s"][si])
|
|
si += 1
|
|
hw = {}
|
|
# Step 1: each input's channel count, then its size.
|
|
for i, t in zip(g["inputs"], allin):
|
|
if t.c != g["tensors"][i["tensor"]]["c"] or t.h < 1 or t.w < 1:
|
|
raise Refuse("model_mismatch")
|
|
if i["max_h"] and (t.h > i["max_h"] or t.w > i["max_w"]):
|
|
raise Refuse("too_large")
|
|
if t.c * t.h * t.w > g["max_elems"]:
|
|
raise Refuse("too_large")
|
|
hw[i["tensor"]] = (t.h, t.w)
|
|
# Step 2: each scalar against its range.
|
|
if len(scalars) != len(g["scalars"]):
|
|
raise Refuse("model_mismatch")
|
|
for v, (lo, hi) in zip(scalars, g["scalars"]):
|
|
if not lo <= v <= hi:
|
|
raise Refuse("field_range")
|
|
# Step 3: the run-time plan.
|
|
walk(g, hw, scalars, False)
|
|
# Step 4: evaluation, the inputs' policies first.
|
|
vals = {}
|
|
for i, t in zip(g["inputs"], allin):
|
|
t = T(t.c, t.h, t.w, list(t.d))
|
|
if not apply_policy(g["tensors"][i["tensor"]], t):
|
|
raise Refuse("field_range")
|
|
vals[i["tensor"]] = t
|
|
for op in g["ops"]:
|
|
outs = ev(g, op, vals, hw, scalars)
|
|
for t, ten in zip(op["outs"], outs):
|
|
# A value outside i64 is one load step 8 excludes; before the certifier
|
|
# exists, such a decode is refused.
|
|
if any(v is None or not I64_MIN <= v <= I64_MAX for v in ten.d):
|
|
raise Refuse("model_mismatch")
|
|
if not apply_policy(g["tensors"][t], ten):
|
|
raise Refuse("field_range")
|
|
vals[t] = ten
|
|
outs = [vals[t] for t in g["outputs"]]
|
|
state["s"] = [outs[i["src"]] for i in g["inputs"] if i["kind"] == 1]
|
|
return outs
|
|
|
|
|
|
# --------------------------------------------------------------- canonical JSON (8.5, 14.1)
|
|
|
|
|
|
def cw(v):
|
|
"""The canonical form: ints as decimal strings, members sorted by key bytes."""
|
|
if v is None:
|
|
return "null"
|
|
if v is True:
|
|
return "true"
|
|
if v is False:
|
|
return "false"
|
|
if isinstance(v, int):
|
|
return '"%d"' % v
|
|
if isinstance(v, str):
|
|
return '"' + v.replace("\\", "\\\\").replace('"', '\\"') + '"'
|
|
if isinstance(v, list):
|
|
return "[" + ",".join(cw(x) for x in v) + "]"
|
|
if isinstance(v, dict):
|
|
keys = sorted(v, key=lambda k: k.encode())
|
|
return "{" + ",".join(cw(k) + ":" + cw(v[k]) for k in keys) + "}"
|
|
raise TypeError(v)
|
|
|
|
|
|
def dec(s):
|
|
"""A canonical decimal integer string."""
|
|
if not isinstance(s, str) or not s or s == "-0":
|
|
raise ValueError("not a canonical integer: %r" % (s,))
|
|
d = s[1:] if s[0] == "-" else s
|
|
if not d.isdigit() or (len(d) > 1 and d[0] == "0") or not d.isascii():
|
|
raise ValueError("not a canonical integer: %r" % (s,))
|
|
return int(s)
|
|
|
|
|
|
def tensor_in(v):
|
|
c, h, w = (dec(x) for x in v["shape"])
|
|
d = [dec(x) for x in v["data"]]
|
|
if len(d) != c * h * w:
|
|
raise ValueError("a tensor's data does not fill its shape")
|
|
return T(c, h, w, d)
|
|
|
|
|
|
def tensor_out(t):
|
|
return {"data": list(t.d), "shape": [t.c, t.h, t.w]}
|
|
|
|
|
|
def graph_line(v):
|
|
"""A graph vector, recomputed: `load` (steps 1-8), `profiles` (section 14.1: both
|
|
when the load is refused, else exactly the profiles the verdict meets) and every
|
|
decode's `expect`."""
|
|
b = bytes.fromhex(v["graph"])
|
|
out = {"graph": b.hex(), "id": v["id"], "decodes": []}
|
|
try:
|
|
g = load(b)
|
|
out["load"] = "ok"
|
|
out["profiles"] = [p for p in ("acc32", "acc64") if g["verdict"][p]]
|
|
except Refuse as e:
|
|
out["load"] = e.code
|
|
out["profiles"] = ["acc32", "acc64"]
|
|
g = None
|
|
if g is not None:
|
|
state = {}
|
|
for d in v["decodes"]:
|
|
ins = [tensor_in(t) for t in d["inputs"]]
|
|
sc = [dec(s) for s in d["scalars"]]
|
|
try:
|
|
res = run(g, ins, sc, state)
|
|
exp = {"outputs": [tensor_out(t) for t in res]}
|
|
except Refuse as e:
|
|
exp = {"refuse": e.code}
|
|
out["decodes"].append({"expect": exp, "inputs": [tensor_out(t) for t in ins], "scalars": sc})
|
|
return out
|
|
|
|
|
|
# --------------------------------------------------------------- stage operators (5E)
|
|
|
|
|
|
class Rans:
|
|
"""integer-payloads section 2's decoder state."""
|
|
|
|
def __init__(self, b):
|
|
self.b = b
|
|
self.x = (int.from_bytes(b[4:8], "little") << 32) | int.from_bytes(b[0:4], "little")
|
|
self.i = 2
|
|
self.out = False
|
|
|
|
def renorm(self):
|
|
if self.x < (1 << 31):
|
|
if 4 * self.i + 4 <= len(self.b):
|
|
self.x = (self.x << 32) | int.from_bytes(self.b[4 * self.i:4 * self.i + 4], "little")
|
|
self.i += 1
|
|
else:
|
|
self.out = True
|
|
|
|
def bits(self, n):
|
|
v = self.x & ((1 << n) - 1)
|
|
self.x >>= n
|
|
self.renorm()
|
|
return v
|
|
|
|
|
|
def sat32(v):
|
|
return max(I32_MIN, min(I32_MAX, v))
|
|
|
|
|
|
def rans64_decode(i, m):
|
|
table_rule, escape, end = m["table"], m["escape"], m["end"]
|
|
lo = hi = None
|
|
if escape == "strict":
|
|
lo, hi = dec(m["lo"]), dec(m["hi"])
|
|
if not (I32_MIN <= lo <= hi <= I32_MAX):
|
|
raise Refuse("model_mismatch", -1)
|
|
tables = [([dec(c) for c in t["cdf"]], dec(t["offset"])) for t in i["tables"]]
|
|
reads = [dec(r) for r in i["reads"]]
|
|
for cdf, o in tables:
|
|
n = len(cdf)
|
|
ok_n = 3 <= n <= 4097 if table_rule == "t4097" else 2 <= n <= 65537
|
|
if not ok_n or cdf[0] != 0 or cdf[-1] != 65536 or any(cdf[k] >= cdf[k + 1] for k in range(n - 1)) or not I32_MIN <= o <= I32_MAX:
|
|
raise Refuse("model_mismatch", -1)
|
|
if any(not 0 <= r < len(tables) for r in reads):
|
|
raise Refuse("model_mismatch", -1)
|
|
stream = bytes.fromhex(i["stream"])
|
|
if len(stream) < 8 or len(stream) % 4:
|
|
raise Refuse("stream_exhausted", 0)
|
|
d = Rans(stream)
|
|
syms = []
|
|
for k, r in enumerate(reads):
|
|
cdf, o = tables[r]
|
|
n = len(cdf)
|
|
cum = d.x & 0xFFFF
|
|
s = max(j for j in range(n - 1) if cdf[j] <= cum)
|
|
d.x = (cdf[s + 1] - cdf[s]) * (d.x >> 16) + cum - cdf[s]
|
|
d.renorm()
|
|
code = None
|
|
if s < n - 2:
|
|
sym = s + o
|
|
if escape == "lenient":
|
|
sym = sat32(sym)
|
|
elif not lo <= sym <= hi:
|
|
code = "field_range"
|
|
elif escape == "lenient":
|
|
v = d.bits(4)
|
|
count = v
|
|
groups = 0
|
|
while v == 15 and groups < 8 and not d.out:
|
|
v = d.bits(4)
|
|
count += v
|
|
groups += 1
|
|
count = min(count, 8)
|
|
raw = 0
|
|
for j in range(count):
|
|
raw |= d.bits(4) << (4 * j)
|
|
value = -(raw >> 1) - 1 if raw & 1 else (raw >> 1) + (n - 2)
|
|
sym = sat32(sat32(value) + o)
|
|
else:
|
|
count = d.bits(4)
|
|
raw, last = 0, None
|
|
if count <= 8:
|
|
for j in range(count):
|
|
last = d.bits(4)
|
|
raw += last << (4 * j)
|
|
value = -(raw // 2) - 1 if raw % 2 else raw // 2 + (n - 2)
|
|
sym = value + o
|
|
if count > 8 or (count >= 1 and last == 0):
|
|
code = "stream_noncanonical"
|
|
elif not lo <= sym <= hi:
|
|
code = "field_range"
|
|
if end == "canonical" and d.out:
|
|
raise Refuse("stream_exhausted", k)
|
|
if code:
|
|
raise Refuse(code, k)
|
|
syms.append(sym)
|
|
if end == "canonical":
|
|
if d.out or d.x != (1 << 31) or 4 * d.i != len(stream):
|
|
raise Refuse("stream_noncanonical", len(reads))
|
|
elif d.out:
|
|
raise Refuse("stream_exhausted", len(reads))
|
|
return {"symbols": syms}
|
|
|
|
|
|
class Win:
|
|
"""E8's window and states over memory M (8 zero bytes, then the stream)."""
|
|
|
|
def __init__(self, stream):
|
|
self.M = bytes(8) + stream
|
|
self.p = len(stream)
|
|
self.h = self.le64(self.p)
|
|
self.pos = 0
|
|
self.last_in = 0
|
|
if self.h == 0:
|
|
raise Refuse("stream_exhausted", 0)
|
|
self.b = self.h.bit_length() - 1
|
|
self.h -= 1 << self.b
|
|
self.s = [0, 0]
|
|
try:
|
|
self.s[0] = self.read(8)
|
|
self.s[1] = self.read(8)
|
|
self.flush()
|
|
except Refuse:
|
|
raise Refuse("stream_exhausted", 0)
|
|
|
|
def le64(self, i):
|
|
return int.from_bytes(self.M[i:i + 8], "little")
|
|
|
|
def read(self, k):
|
|
if k > self.b:
|
|
raise Refuse("stream_exhausted", self.pos)
|
|
self.b -= k
|
|
r = self.h >> self.b
|
|
self.h -= r << self.b
|
|
return r
|
|
|
|
def flush(self):
|
|
d = 7 ^ (self.b // 8)
|
|
if d > self.p:
|
|
raise Refuse("stream_exhausted", self.pos)
|
|
self.p -= d
|
|
self.h = self.le64(self.p) % (1 << (56 + self.b % 8))
|
|
self.b = 56 + self.b % 8
|
|
|
|
def step(self, which, t):
|
|
self.s[which] = ((t >> 16) | self.read(t >> 24)) % 256
|
|
|
|
def at(self, k):
|
|
"""The `in` or `out` step of position k runs: a refusal in it is at k."""
|
|
self.pos = k
|
|
|
|
def at_in(self, k):
|
|
self.pos = k
|
|
self.last_in = k
|
|
|
|
def group_flush(self):
|
|
"""A flush that belongs to no single position (section 14.1, `at`): it reports
|
|
the last position whose `in` step ran."""
|
|
self.pos = self.last_in
|
|
self.flush()
|
|
|
|
|
|
def wrap16(a):
|
|
return ((a + 32768) % 65536) - 32768
|
|
|
|
|
|
def latent_tables(pmf):
|
|
w = len(pmf[0]) if pmf else 0
|
|
if w < 1 or any(len(r) != w or min(r) < 0 or sum(r) != 256 for r in pmf):
|
|
raise Refuse("model_mismatch", -1)
|
|
tabs = []
|
|
for P in pmf:
|
|
C = [0]
|
|
for p in P:
|
|
C.append(C[-1] + p)
|
|
F = lambda k: C[k] - C[k] // 2
|
|
G = lambda k: C[k] // 2 + 128
|
|
sym = lambda k: k // 2 if k % 2 == 0 else -((k + 1) // 2)
|
|
slot_sym, slot_next = [None] * 256, [None] * 256
|
|
k = 0
|
|
while C[k] < 256:
|
|
for u in range(F(k), F(k + 1)):
|
|
slot_sym[u], slot_next[u] = sym(k), P[k] + (u - F(k))
|
|
for u in range(G(k), G(k + 1)):
|
|
slot_sym[u], slot_next[u] = sym(k), P[k] + (F(k + 1) - F(k)) + (u - G(k))
|
|
k += 1
|
|
t = []
|
|
for u in range(256):
|
|
nb = 0
|
|
while slot_next[u] * (1 << nb) < 256:
|
|
nb += 1
|
|
base = (slot_next[u] * (1 << nb)) ^ 256
|
|
t.append(nb * (1 << 24) + base * (1 << 16) + slot_sym[u] % 65536)
|
|
tabs.append(t)
|
|
return tabs
|
|
|
|
|
|
def mans_sgm(i):
|
|
pmf = [[dec(x) for x in r] for r in i["pmf"]]
|
|
bounds = [dec(x) for x in i["bounds"]]
|
|
spans = [([bool(b) for b in s["mask"]], [dec(x) for x in s["sigma"]]) for s in i["spans"]]
|
|
tabs = latent_tables(pmf)
|
|
if any(not 0 <= b <= 255 for b in bounds):
|
|
raise Refuse("model_mismatch", -1)
|
|
nd = min(len(tabs), len(bounds))
|
|
if any(len(m) != len(sg) or any(not 0 <= x < nd for x in sg) for m, sg in spans):
|
|
raise Refuse("model_mismatch", -1)
|
|
w = Win(bytes.fromhex(i["memory"]))
|
|
out = []
|
|
base = 0
|
|
for mask, sigma in spans:
|
|
L = len(mask)
|
|
sy = [0] * L
|
|
|
|
def inb(which, k):
|
|
w.at_in(base + k)
|
|
if mask[k]:
|
|
t = tabs[sigma[k]][w.s[which]]
|
|
w.step(which, t)
|
|
sy[k] = wrap16(t % 65536)
|
|
|
|
def outb(k):
|
|
w.at(base + k)
|
|
if mask[k]:
|
|
beta = bounds[sigma[k]]
|
|
if sy[k] + beta == 0:
|
|
f = w.read(1)
|
|
v = w.read(3 if f == 1 else 16)
|
|
sy[k] = wrap16(v // 2 + beta) if v % 2 == 0 else wrap16(-(v // 2) - beta)
|
|
w.flush()
|
|
|
|
k = 0
|
|
while k + 3 < L:
|
|
inb(0, k); inb(1, k + 1); inb(0, k + 2); inb(1, k + 3)
|
|
w.group_flush()
|
|
for q in range(4):
|
|
outb(k + q)
|
|
k += 4
|
|
if L % 4 >= 2:
|
|
inb(0, k); outb(k); inb(1, k + 1); outb(k + 1)
|
|
k += 2
|
|
w.group_flush()
|
|
if L % 2 == 1:
|
|
inb(0, k); outb(k)
|
|
w.group_flush()
|
|
out.append(sy)
|
|
base += L
|
|
return {"symbols": out}
|
|
|
|
|
|
def mans_fac(i):
|
|
rows = [[dec(x) for x in r] for r in i["rows"]]
|
|
lens = [dec(x) for x in i["lengths"]]
|
|
if len(rows) != len(lens) or any(len(r) != 63 or r[0] < 0 or r[62] > 255 or any(r[k] > r[k + 1] for k in range(62)) for r in rows):
|
|
raise Refuse("model_mismatch", -1)
|
|
Zp = [0] * 512
|
|
nb = 0
|
|
for i_ in range(511, 0, -1):
|
|
if i_ * (1 << nb) < 256:
|
|
nb += 1
|
|
Zp[i_] = ((i_ * (1 << nb)) ^ 256) + 256 * nb
|
|
w = Win(bytes.fromhex(i["memory"]))
|
|
Tz = [0] * 256
|
|
out = []
|
|
base = 0
|
|
for r, L in zip(rows, lens):
|
|
prev = 0
|
|
for a in range(63):
|
|
k = r[a] - prev
|
|
for j in range(prev, r[a]):
|
|
Tz[j] = Zp[k + j - prev] * 65536 + a
|
|
prev = r[a]
|
|
Tz[255] = Zp[1] * 65536 + 63
|
|
sy = [0] * L
|
|
|
|
def inb(which, k):
|
|
w.at_in(base + k)
|
|
t = Tz[w.s[which]]
|
|
w.step(which, t)
|
|
sy[k] = t % 256
|
|
|
|
def outb(k):
|
|
w.at(base + k)
|
|
if sy[k] == 63:
|
|
sy[k] = w.read(6)
|
|
|
|
k = 0
|
|
while k + 3 < L:
|
|
inb(0, k); inb(1, k + 1); inb(0, k + 2); inb(1, k + 3)
|
|
for q in range(4):
|
|
outb(k + q)
|
|
w.group_flush()
|
|
k += 4
|
|
if L % 4 >= 2:
|
|
inb(0, k); outb(k); inb(1, k + 1); outb(k + 1)
|
|
k += 2
|
|
if L % 2 == 1:
|
|
inb(0, k); outb(k)
|
|
w.group_flush()
|
|
out.append(sy)
|
|
base += L
|
|
return {"symbols": out}
|
|
|
|
|
|
def rdiv(n, d):
|
|
if d < 1:
|
|
return 0
|
|
sg = (n > 0) - (n < 0)
|
|
return sg * ((2 * abs(n) + d) // (2 * d))
|
|
|
|
|
|
def stage_expect(op, i, m):
|
|
D = lambda k: dec(i[k])
|
|
if op == "rans64_decode":
|
|
return rans64_decode(i, m)
|
|
if op == "dequantize":
|
|
f, s, b = D("f"), D("sym"), D("base")
|
|
if not 0 <= f <= 62:
|
|
raise Refuse("field_range", 0)
|
|
v = s * (1 << f) + b
|
|
if not I64_MIN <= v <= I64_MAX:
|
|
raise Refuse("field_range", 0)
|
|
return {"value": v}
|
|
if op == "bucket":
|
|
sigma, e = D("sigma"), [dec(x) for x in i["edges"]]
|
|
if m["bucket"] == "ceil":
|
|
if not e:
|
|
raise Refuse("model_mismatch", -1)
|
|
s2 = max(sigma, e[0])
|
|
return {"index": (len(e) - 1) - sum(1 for k in range(len(e) - 1) if s2 <= e[k])}
|
|
if not e or any(e[k] > e[k + 1] for k in range(len(e) - 1)):
|
|
raise Refuse("model_mismatch", -1)
|
|
return {"index": sum(1 for x in e if x <= sigma)}
|
|
if op == "coded":
|
|
return {"coded": D("v") > D("thr")}
|
|
if op == "scale_index":
|
|
n, s = D("n"), D("s")
|
|
if n < 1 or not 0 <= s <= 127:
|
|
raise Refuse("model_mismatch", -1)
|
|
return {"index": min(max(R(D("v"), s, 1), 0), n - 1)}
|
|
if op == "block_stat":
|
|
x = tensor_in(i["x"])
|
|
b, s, lo, hi = D("b"), D("s"), D("lo"), D("hi")
|
|
if b < 1 or not 0 <= s <= 127 or lo > hi:
|
|
raise Refuse("model_mismatch", -1)
|
|
d = list(x.d)
|
|
for c in range(x.c):
|
|
for by in range(0, x.h, b):
|
|
for bx in range(0, x.w, b):
|
|
cells = [(y, xx) for y in range(by, min(by + b, x.h)) for xx in range(bx, min(bx + b, x.w))]
|
|
v = min(max(R(sum(x.at(c, y, xx) for y, xx in cells), s, 1), lo), hi)
|
|
for y, xx in cells:
|
|
d[(c * x.h + y) * x.w + xx] = v
|
|
return {"out": tensor_out(T(x.c, x.h, x.w, d))}
|
|
if op == "band_lookup":
|
|
e = [dec(x) for x in i["edges"]]
|
|
table = [[dec(x) for x in r] for r in i["table"]]
|
|
row = D("row")
|
|
if not 0 <= row < len(table) or len(table[row]) != len(e) + 1:
|
|
raise Refuse("model_mismatch", -1)
|
|
return {"value": table[row][sum(1 for k in e if D("v") >= k)]}
|
|
if op == "mans_decode":
|
|
return mans_sgm(i) if i["mode"] == "sgm" else mans_fac(i)
|
|
if op == "residual_step":
|
|
pred, r, q, lo, hi = D("pred"), D("r"), D("q"), D("lo"), D("hi")
|
|
if lo >= hi or not lo <= pred <= hi or not 1 <= q <= 65535 or not I32_MIN <= r <= I32_MAX:
|
|
raise Refuse("model_mismatch", -1)
|
|
return {"value": min(max(pred + r * q, lo), hi)}
|
|
if op == "mc_predict":
|
|
f = tensor_in(i["frame"])
|
|
y, x, c, dy, dx, pr = D("y"), D("x"), D("c"), D("dy"), D("dx"), D("precision")
|
|
if not (0 <= c < f.c and 0 <= y < f.h and 0 <= x < f.w) or pr not in (1, 2):
|
|
raise Refuse("model_mismatch", -1)
|
|
p = lambda Y, X: f.at(c, min(max(Y, 0), f.h - 1), min(max(X, 0), f.w - 1))
|
|
if pr == 1:
|
|
return {"value": p(y + dy, x + dx)}
|
|
iy, fy, ix, fx = y + (dy >> 1), dy & 1, x + (dx >> 1), dx & 1
|
|
if not fy and not fx:
|
|
v = p(iy, ix)
|
|
elif not fy:
|
|
v = (p(iy, ix) + p(iy, ix + 1) + 1) >> 1
|
|
elif not fx:
|
|
v = (p(iy, ix) + p(iy + 1, ix) + 1) >> 1
|
|
else:
|
|
v = (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) >> 2
|
|
return {"value": v}
|
|
if op == "up2_predict":
|
|
prev = tensor_in(i["prev"])
|
|
d0, d1 = D("d0"), D("d1")
|
|
hp, wp = -(-d0 // 2), -(-d1 // 2)
|
|
if d0 < 1 or d1 < 1 or (prev.h, prev.w) != (hp, wp):
|
|
raise Refuse("geometry_mismatch", -1)
|
|
p = lambda c, Y, X: prev.at(c, min(Y, hp - 1), min(X, wp - 1))
|
|
d = []
|
|
for c in range(prev.c):
|
|
for y in range(d0):
|
|
for x in range(d1):
|
|
iy, fy, ix, fx = y >> 1, y & 1, x >> 1, x & 1
|
|
if not fy and not fx:
|
|
v = p(c, iy, ix)
|
|
elif not fy:
|
|
v = (p(c, iy, ix) + p(c, iy, ix + 1) + 1) >> 1
|
|
elif not fx:
|
|
v = (p(c, iy, ix) + p(c, iy + 1, ix) + 1) >> 1
|
|
else:
|
|
v = (p(c, iy, ix) + p(c, iy, ix + 1) + p(c, iy + 1, ix) + p(c, iy + 1, ix + 1) + 2) >> 2
|
|
d.append(v)
|
|
return {"out": tensor_out(T(prev.c, d0, d1, d))}
|
|
if op == "rate_predict":
|
|
prev = tensor_in(i["prev"])
|
|
num, den, d0 = D("rate_num"), D("rate_den"), D("d0")
|
|
n = prev.w
|
|
if prev.h != 1 or num < 1 or den < 1 or n * num // den != d0:
|
|
raise Refuse("geometry_mismatch", -1)
|
|
d = []
|
|
for c in range(prev.c):
|
|
for k in range(d0):
|
|
s = k * den
|
|
lo, rho = s // num, s % num
|
|
hi = min(lo + 1, n - 1)
|
|
a, b = prev.at(c, 0, lo), prev.at(c, 0, hi)
|
|
d.append(rdiv(a * num + (b - a) * rho, num))
|
|
return {"out": tensor_out(T(prev.c, 1, d0, d))}
|
|
if op == "coding_order":
|
|
c, h, w = (dec(x) for x in i["shape"])
|
|
cb = None if i["order"] == "raster" else i["anchor"]
|
|
idx = [(cc * h + y) * w + x for cc in range(c) for y in range(h) for x in range(w) if cb is None or (((y + x) % 2 == 0) == cb)]
|
|
return {"indices": idx}
|
|
raise ValueError("unknown stage operator %r" % op)
|
|
|
|
|
|
def stage_line(v):
|
|
try:
|
|
exp = stage_expect(v["op"], v["in"], v["modes"])
|
|
except Refuse as e:
|
|
exp = {"at": e.at, "refuse": e.code}
|
|
return {"expect": exp, "id": v["id"], "in": v["in"], "modes": v["modes"], "op": v["op"]}
|
|
|
|
|
|
# --------------------------------------------------------------- lowerings (section 13)
|
|
|
|
|
|
class WB:
|
|
"""A WHG1 graph under construction with section 13.0's numbering."""
|
|
|
|
def __init__(self):
|
|
self.scalars, self.inputs, self.tensors, self.ops = [], [], [], []
|
|
|
|
def tensor(self, c, policy=0, lo=0, hi=0):
|
|
self.tensors.append((c, policy, lo, hi))
|
|
return len(self.tensors) - 1
|
|
|
|
def input(self, c, lo, hi, mh=0, mw=0):
|
|
t = self.tensor(c, 2, lo, hi)
|
|
self.inputs.append((t, mh, mw))
|
|
return t
|
|
|
|
def op(self, code, ins, out_rec, params):
|
|
t = self.tensor(*out_rec)
|
|
self.ops.append((code, ins, [t], params))
|
|
return t
|
|
|
|
def bytes(self, outputs):
|
|
b = bytearray(b"WHG1")
|
|
b += bytes([1, 1, 0, 0])
|
|
b += struct.pack("<IIHHHHII", len(self.tensors), len(self.ops), len(self.inputs), len(outputs), len(self.scalars), 0, CAP, 0)
|
|
for lo, hi in self.scalars:
|
|
b += struct.pack("<qq", lo, hi)
|
|
for t, mh, mw in self.inputs:
|
|
b += struct.pack("<IIIB3xI4xq", t, mh, mw, 0, 0, 0)
|
|
for t in outputs:
|
|
b += struct.pack("<I", t)
|
|
for c, pol, lo, hi in self.tensors:
|
|
b += struct.pack("<Ibb2xqq", c, -128, pol, lo, hi)
|
|
for code, ins, outs, params in self.ops:
|
|
b += struct.pack("<BBBx", code, len(ins), len(outs))
|
|
for t in ins + outs:
|
|
b += struct.pack("<I", t)
|
|
b += struct.pack("<I", len(params)) + params
|
|
return bytes(b)
|
|
|
|
|
|
def ints_le(v, width):
|
|
return b"".join(x.to_bytes(width, "little", signed=True) for x in v)
|
|
|
|
|
|
def rq_block(round_, bias, mult, mkind, mbytes, shifts, skind, biasv, bbytes):
|
|
"""The RQ block of section 10.4."""
|
|
b = struct.pack("<BBBBBBH", round_, bias, mkind, skind, mbytes, bbytes, 0)
|
|
b += ints_le(mult, mbytes) if mkind else b""
|
|
b += bytes(shifts)
|
|
b += ints_le(biasv, bbytes) if bias else b""
|
|
return b
|
|
|
|
|
|
def smallest(v, allowed):
|
|
for w in allowed:
|
|
lo, hi = -(1 << (8 * w - 1)), (1 << (8 * w - 1)) - 1
|
|
if all(lo <= x <= hi for x in v):
|
|
return w
|
|
return 8
|
|
|
|
|
|
class LowerRefuse(Exception):
|
|
pass
|
|
|
|
|
|
def lref():
|
|
raise LowerRefuse()
|
|
|
|
|
|
def p_conv(cout, kh, kw, sh, sw, ph, pw, groups, w_bytes, rq, w):
|
|
return struct.pack("<8IB3x", cout, kh, kw, sh, sw, ph, pw, groups, w_bytes) + rq + ints_le(w, w_bytes)
|
|
|
|
|
|
def p_convt(cout, kh, kw, sh, sw, ph, pw, oph, opw, w_bytes, rq, w):
|
|
return struct.pack("<9IB3x", cout, kh, kw, sh, sw, ph, pw, oph, opw, w_bytes) + rq + ints_le(w, w_bytes)
|
|
|
|
|
|
def p_clamp(lo, hi):
|
|
return struct.pack("<qq", lo, hi)
|
|
|
|
|
|
def int_synth_lower(js):
|
|
"""int_synth/1 (section 13.1)."""
|
|
try:
|
|
root = json.loads(js)
|
|
except ValueError:
|
|
lref()
|
|
layers = root.get("layers") if isinstance(root, dict) else None
|
|
if not isinstance(layers, list) or not layers:
|
|
lref()
|
|
|
|
def u(l, k, default=None):
|
|
v = l.get(k, default)
|
|
if v is None or isinstance(v, bool) or not isinstance(v, int) or not 0 <= v <= 0xFFFFFFFF:
|
|
lref()
|
|
return v
|
|
|
|
def i32s(l, k):
|
|
a = l.get(k)
|
|
if not isinstance(a, list) or any(isinstance(x, bool) or not isinstance(x, int) or not I32_MIN <= x <= I32_MAX for x in a):
|
|
lref()
|
|
return a
|
|
|
|
if not isinstance(layers[0], dict):
|
|
lref()
|
|
c = u(layers[0], "in_ch")
|
|
if c == 0:
|
|
lref()
|
|
wb = WB()
|
|
x = wb.input(c, -128, 127)
|
|
for l in layers:
|
|
if not isinstance(l, dict):
|
|
lref()
|
|
ic, oc, k, s, p = u(l, "in_ch"), u(l, "out_ch"), u(l, "k"), u(l, "stride"), u(l, "pad")
|
|
op_, sh = u(l, "out_pad", 0), u(l, "shift", 0)
|
|
relu = l.get("relu", False)
|
|
if not isinstance(relu, bool):
|
|
lref()
|
|
w, b = i32s(l, "weight"), i32s(l, "bias")
|
|
if ic != c or oc == 0 or k == 0 or s == 0 or sh > 127 or len(w) != ic * oc * k * k or len(b) != oc:
|
|
lref()
|
|
rq = rq_block(0, 1, [], 0, 0, [sh], 1, b, 4)
|
|
x = wb.op(0x11, [x], (oc,), p_convt(oc, k, k, s, s, p, p, op_, op_, 4, rq, w))
|
|
if relu:
|
|
x = wb.op(0x0C, [x], (oc,), b"")
|
|
c = oc
|
|
x = wb.op(0x0B, [x], (c,), p_clamp(0, 255))
|
|
return [("synth", wb.bytes([x]))]
|
|
|
|
|
|
def parse_net(b, magic, audio):
|
|
"""WHM1 / WAM1 (integer-payloads section 9)."""
|
|
if b[:4] != magic:
|
|
lref()
|
|
rd = Reader(b, 4)
|
|
try:
|
|
afrac, gf, shift, n = rd.u(4), rd.u(4), rd.u(4), rd.u(4)
|
|
layers = []
|
|
for _ in range(n):
|
|
tag = rd.u(1)
|
|
if tag == 0:
|
|
inc, out, k, s, p, op_ = (rd.u(4) for _ in range(6))
|
|
w = rd.ints(inc * out * k * (1 if audio else k), 2)
|
|
layers.append(("convT", inc, out, k, s, p, op_, w, rd.ints(out, 8), rd.ints(out, 8)))
|
|
elif tag == 1 and not audio:
|
|
inc, out, k, s, p = (rd.u(4) for _ in range(5))
|
|
w = rd.ints(out * inc * k * k, 2)
|
|
layers.append(("conv", inc, out, k, s, p, 0, w, rd.ints(out, 8), rd.ints(out, 8)))
|
|
elif tag == 2:
|
|
c = rd.u(4)
|
|
layers.append(("igdn", c, rd.ints(c * c, 4), rd.ints(c, 8)))
|
|
elif tag == 3 and not audio:
|
|
layers.append(("relu",))
|
|
elif tag == 4 and not audio:
|
|
layers.append(("leaky", rd.i(8), rd.u(4)))
|
|
else:
|
|
lref()
|
|
except Refuse:
|
|
lref()
|
|
if rd.p != len(b):
|
|
lref()
|
|
return {"afrac": afrac, "gf": gf, "shift": shift, "layers": layers, "audio": audio}
|
|
|
|
|
|
def largest_x(pred):
|
|
"""The largest X in [0, 2^63] with pred(X), pred monotone; None if pred(0) fails."""
|
|
if not pred(0):
|
|
return None
|
|
lo, hi = 0, 1 << 63
|
|
if pred(hi):
|
|
return hi
|
|
while hi - lo > 1:
|
|
mid = (lo + hi) // 2
|
|
if pred(mid):
|
|
lo = mid
|
|
else:
|
|
hi = mid
|
|
return lo
|
|
|
|
|
|
def x_l(net, layer, c):
|
|
"""X_l of section 13.2 for a layer of tag 0, 1, 2 or 4 reading c channels; the set
|
|
is refused when P_l(0) fails."""
|
|
I64, I128 = (1 << 63) - 1, (1 << 127) - 1
|
|
kind = layer[0]
|
|
if kind in ("conv", "convT"):
|
|
k = layer[3]
|
|
terms = c * k * (1 if net["audio"] else k)
|
|
Wm, Mm, Bm = max(map(abs, layer[7]), default=0), max(map(abs, layer[8]), default=0), max(map(abs, layer[9]), default=0)
|
|
p = lambda X: terms * Wm * X <= I64 and (terms * Wm * X * Mm) // (1 << net["shift"]) + 1 + Bm <= I64
|
|
elif kind == "igdn":
|
|
G, Bb = max(map(abs, layer[2]), default=0), max(map(abs, layer[3]), default=0)
|
|
|
|
def p(X):
|
|
if G * X * X > I128:
|
|
return False
|
|
N = Bb + c * ((G * X * X) // (1 << net["gf"]) + 1)
|
|
if N > I128:
|
|
return False
|
|
r = math.isqrt(N)
|
|
return r <= I64 and (X * r) // (1 << net["afrac"]) + 1 <= I64
|
|
else:
|
|
a = abs(layer[1])
|
|
p = lambda X: X * a <= I64
|
|
x = largest_x(p)
|
|
if x is None:
|
|
lref()
|
|
return x
|
|
|
|
|
|
def x_record(net, k, c):
|
|
"""The tensor record of x_k on c channels (section 13.2, Range checks). x_k is read
|
|
by layer k, if there is one. When that layer is of tag 0, 1, 2 or 4, x_k is
|
|
check [-X_k, min(X_k, I64)]; that includes the network's input. A network input
|
|
that no such layer reads is check [-2^63, I64]. Every other tensor is derived."""
|
|
L = net["layers"]
|
|
if k < len(L) and L[k][0] in ("convT", "conv", "igdn", "leaky"):
|
|
x = x_l(net, L[k], c)
|
|
return (c, 2, -x, min(x, I64_MAX))
|
|
if k == 0:
|
|
return (c, 2, I64_MIN, I64_MAX)
|
|
return (c,)
|
|
|
|
|
|
def net_layers(wb, net, x):
|
|
L = net["layers"]
|
|
for k, l in enumerate(L):
|
|
c = wb.tensors[x][0]
|
|
oc = l[2] if l[0] in ("conv", "convT") else c
|
|
rec = x_record(net, k + 1, oc)
|
|
if l[0] in ("conv", "convT"):
|
|
_, inc, out, kk, s, p, op_, w, mult, bias = l
|
|
if inc != c or out == 0 or kk == 0 or s == 0:
|
|
lref()
|
|
rq = rq_block(0, 2, mult, 2, 8, [net["shift"]], 1, bias, 8)
|
|
if l[0] == "conv":
|
|
x = wb.op(0x01, [x], rec, p_conv(out, kk, kk, s, s, p, p, 1, 2, rq, w))
|
|
elif net["audio"]:
|
|
x = wb.op(0x11, [x], rec, p_convt(out, 1, kk, 1, s, 0, p, 0, op_, 2, rq, w))
|
|
else:
|
|
x = wb.op(0x11, [x], rec, p_convt(out, kk, kk, s, s, p, p, op_, op_, 2, rq, w))
|
|
elif l[0] == "igdn":
|
|
if l[1] != c:
|
|
lref()
|
|
x = wb.op(0x12, [x], rec, struct.pack("<IIB7x", net["gf"], net["afrac"], 4) + ints_le(l[2], 4) + ints_le(l[3], 8))
|
|
elif l[0] == "relu":
|
|
x = wb.op(0x0C, [x], rec, b"")
|
|
else:
|
|
if l[2] > 63:
|
|
lref()
|
|
x = wb.op(0x0D, [x], rec, struct.pack("<qI", l[1], l[2]))
|
|
return x
|
|
|
|
|
|
def net_check(net):
|
|
if net["shift"] > 127 or net["gf"] > 127 or not 1 <= net["afrac"] <= 30 or not net["layers"]:
|
|
lref()
|
|
for l in net["layers"]:
|
|
if l[0] in ("conv", "convT"):
|
|
return l[1]
|
|
if l[0] == "igdn":
|
|
return l[1]
|
|
lref()
|
|
|
|
|
|
def rgb8(wb, x, afrac):
|
|
c = wb.tensors[x][0]
|
|
t = wb.op(0x0B, [x], (c,), p_clamp(0, 1 << afrac))
|
|
return wb.op(0x02, [t], (c,), rq_block(1, 0, [255], 1, 4, [afrac], 1, [], 0))
|
|
|
|
|
|
def whm1_lower(b, output_step):
|
|
net = parse_net(b, b"WHM1", False)
|
|
c0 = net_check(net)
|
|
wb = WB()
|
|
r = x_record(net, 0, c0)
|
|
x = wb.input(r[0], r[2], r[3])
|
|
x = net_layers(wb, net, x)
|
|
if output_step:
|
|
x = rgb8(wb, x, net["afrac"])
|
|
return wb.bytes([x])
|
|
|
|
|
|
def wam1_lower(b):
|
|
net = parse_net(b, b"WAM1", True)
|
|
c0 = net_check(net)
|
|
wb = WB()
|
|
r = x_record(net, 0, c0)
|
|
x = wb.input(r[0], r[2], r[3])
|
|
wb.scalars.append((1, 1 << 26))
|
|
x = net_layers(wb, net, x)
|
|
c = wb.tensors[x][0]
|
|
lo = -((1 << 63) // 32767)
|
|
hi = (I64_MAX - (1 << (net["afrac"] - 1))) // 32767
|
|
crop = struct.pack("<B3xIB3xIB3xIB3xI", 0, 0, 0, 0, 0, 1, 1, 0)
|
|
t = wb.op(0x14, [x], (c, 2, lo, hi), crop)
|
|
t = wb.op(0x02, [t], (c,), rq_block(1, 0, [32767], 1, 4, [net["afrac"]], 1, [], 0))
|
|
t = wb.op(0x0B, [t], (c,), p_clamp(-32768, 32767))
|
|
return wb.bytes([t])
|
|
|
|
|
|
def wmb1_lower(b):
|
|
if b[:4] != b"WMB1":
|
|
lref()
|
|
rd = Reader(b, 4)
|
|
try:
|
|
secs = [rd.take(rd.u(4)) for _ in range(4)]
|
|
except Refuse:
|
|
lref()
|
|
if rd.p != len(b):
|
|
lref()
|
|
nets = [parse_net(s, b"WHM1", False) for s in secs[:3]]
|
|
if len({n["afrac"] for n in nets}) != 1:
|
|
lref()
|
|
try:
|
|
tables = json.loads(secs[3])
|
|
except ValueError:
|
|
lref()
|
|
a = tables.get("afrac") if isinstance(tables, dict) else None
|
|
if isinstance(a, bool) or not isinstance(a, int) or a != nets[2]["afrac"]:
|
|
lref()
|
|
return [("h_s", net_graph(nets[0], False)), ("entropy", net_graph(nets[1], False)), ("g_s", net_graph(nets[2], True))]
|
|
|
|
|
|
def net_graph(net, step):
|
|
c0 = net_check(net)
|
|
wb = WB()
|
|
r = x_record(net, 0, c0)
|
|
x = wb.input(r[0], r[2], r[3])
|
|
x = net_layers(wb, net, x)
|
|
if step:
|
|
x = rgb8(wb, x, net["afrac"])
|
|
return wb.bytes([x])
|
|
|
|
|
|
def world_lower(tb, gb, c):
|
|
tn, gn = parse_net(tb, b"WHM1", False), parse_net(gb, b"WHM1", False)
|
|
cin = net_check(tn)
|
|
net_check(gn)
|
|
if tn["afrac"] != gn["afrac"] or c < 1 or cin <= c:
|
|
lref()
|
|
A = cin - c
|
|
wb = WB()
|
|
z = wb.input(c, I64_MIN, I64_MAX)
|
|
parts = [z]
|
|
for _ in range(A):
|
|
wb.scalars.append((I64_MIN, I64_MAX))
|
|
s = len(wb.scalars) - 1
|
|
parts.append(wb.op(0x17, [z], (1,), struct.pack("<B7xq", 1, s)))
|
|
x0 = wb.op(0x06, parts, x_record(tn, 0, cin), b"")
|
|
x = net_layers(wb, tn, x0)
|
|
one = 1 << tn["afrac"]
|
|
out = wb.op(0x0B, [x], (wb.tensors[x][0],), p_clamp(-one, one))
|
|
return [("T", wb.bytes([out])), ("g_s", net_graph(gn, True))]
|
|
|
|
|
|
def jpegai_lower(norm_b, hy_b, huv_b):
|
|
"""jpegai-int/1, section 13.6.1."""
|
|
try:
|
|
norm, hy, huv = json.loads(norm_b), json.loads(hy_b), json.loads(huv_b)
|
|
except ValueError:
|
|
lref()
|
|
isint = lambda v: isinstance(v, int) and not isinstance(v, bool)
|
|
ints = lambda v: v if isinstance(v, list) and all(isint(x) for x in v) else lref()
|
|
if not isinstance(norm, dict):
|
|
lref()
|
|
pmf = norm.get("pmf")
|
|
if not isinstance(pmf, list) or len(pmf) < 32:
|
|
lref()
|
|
widths = set()
|
|
for row in pmf:
|
|
r = ints(row)
|
|
if not r or min(r) < 0 or sum(r) != 256:
|
|
lref()
|
|
widths.add(len(r))
|
|
if len(widths) != 1:
|
|
lref()
|
|
bounds = ints(norm.get("bounds"))
|
|
if len(bounds) < 32 or any(not 0 <= v <= 255 for v in bounds):
|
|
lref()
|
|
if not isint(norm.get("per_channel")) or norm["per_channel"] != 16:
|
|
lref()
|
|
for k, rows in (("cdfs_y", 160), ("cdfs_uv", 96)):
|
|
v = ints(norm.get(k))
|
|
if len(v) != rows * 63 or any(not 0 <= x <= 255 for x in v) or any(v[r * 63 + j] > v[r * 63 + j + 1] for r in range(rows) for j in range(62)):
|
|
lref()
|
|
gy, guv = ints(norm.get("gain_y")), ints(norm.get("gain_uv"))
|
|
if len(gy) != 160 or len(guv) != 96 or any(not I64_MIN <= v <= I64_MAX for v in gy + guv):
|
|
lref()
|
|
|
|
def member(o, k, dims):
|
|
m = o.get(k)
|
|
if not isinstance(m, dict) or m.get("dims") != dims:
|
|
lref()
|
|
v = ints(m.get("v"))
|
|
n = 1
|
|
for d in dims:
|
|
n *= d
|
|
if len(v) != n or any(not I64_MIN <= x <= I64_MAX for x in v):
|
|
lref()
|
|
return v
|
|
|
|
def plane(o, cz, gains):
|
|
if not isinstance(o, dict):
|
|
lref()
|
|
|
|
def d0(k):
|
|
m = o.get(k)
|
|
dims = m.get("dims") if isinstance(m, dict) else None
|
|
if not isinstance(dims, list) or not dims or not isint(dims[0]) or not 1 <= dims[0] <= 0xFFFFFFFF:
|
|
lref()
|
|
return dims[0]
|
|
|
|
c1, c2 = d0("conv1.weight"), d0("depthwise.weight")
|
|
c3 = 16 * cz
|
|
wb = WB()
|
|
x = wb.input(cz, -31, 32, 4, 4)
|
|
blocks = (("conv1", c1, cz, 1, 0, "onnx::Pow_85"), ("depthwise", c2, c1, 3, 1, "onnx::Pow_91"), ("pointwise", c3, c2, 1, 0, "onnx::Pow_97"))
|
|
for bi, (name, cout, cin, k, pad, sh) in enumerate(blocks):
|
|
w = member(o, name + ".weight", [cout, cin, k, k])
|
|
bias = member(o, name + ".bias", [cout])
|
|
shifts = member(o, sh, [1, cout, 1, 1])
|
|
if any(not 0 <= s <= 127 for s in shifts):
|
|
lref()
|
|
x = wb.op(0x0B, [x], (wb.tensors[x][0],), p_clamp(-128, 127))
|
|
rq = rq_block(0, 1, [], 0, 0, shifts, 2, bias, smallest(bias, (4, 8)))
|
|
wbytes = smallest(w, (1, 2, 4, 8))
|
|
x = wb.op(0x01, [x], (cout,), p_conv(cout, k, k, 1, 1, pad, pad, 1, wbytes, rq, w))
|
|
if bi < 2:
|
|
x = wb.op(0x0C, [x], (cout,), b"")
|
|
x = wb.op(0x09, [x], (cz,), struct.pack("<I", 4))
|
|
x = wb.op(0x13, [x], (cz,), b"")
|
|
x = wb.op(0x0B, [x], (cz,), p_clamp(0, 3967))
|
|
x = wb.op(0x02, [x], (cz,), rq_block(0, 1, [], 0, 0, [0], 1, gains, smallest(gains, (4, 8))))
|
|
return wb.bytes([x])
|
|
|
|
return [("hs_y", plane(hy, 160, gy)), ("hs_uv", plane(huv, 96, guv))]
|
|
|
|
|
|
def lower_set(entry, files):
|
|
low, cap = entry["lowering"], entry["capability"]
|
|
try:
|
|
if low == "int_synth/1":
|
|
roles = int_synth_lower(files["int_synth.json"])
|
|
elif low == "whm1/1" and cap == "wai.splat.int_codec":
|
|
roles = [("g_s", whm1_lower(files["g_s.whm"], True))]
|
|
elif low == "whm1/1":
|
|
roles = [("h_s", whm1_lower(files["h_s.whm"], False)), ("g_s", whm1_lower(files["g_s.whm"], True))]
|
|
elif low == "wam1/1":
|
|
roles = [("g_s", wam1_lower(files["model.wam"]))]
|
|
elif low == "wmb1/1":
|
|
roles = wmb1_lower(files["model.wmb"])
|
|
elif low == "world/1":
|
|
roles = world_lower(files["T.whm"], files["g_s.whm"], dec(entry["c"]))
|
|
elif low == "jpegai-int/1":
|
|
roles = jpegai_lower(files["norm.json"], files["weights/hs_y.json"], files["weights/hs_uv.json"])
|
|
else:
|
|
raise ValueError("no lowering %r" % low)
|
|
except LowerRefuse:
|
|
raise ValueError("the lowering refused the set")
|
|
# Every lowered graph loads.
|
|
for _, b in roles:
|
|
load(b)
|
|
return roles
|
|
|
|
|
|
def set_digest(files):
|
|
if len(files) == 1:
|
|
return hashlib.sha256(next(iter(files.values()))).hexdigest()
|
|
listing = "".join("%s %s\n" % (hashlib.sha256(files[n]).hexdigest(), n) for n in sorted(files, key=lambda s: s.encode()))
|
|
return hashlib.sha256(listing.encode()).hexdigest()
|
|
|
|
|
|
# --------------------------------------------------------------- main
|
|
|
|
|
|
def check(root, write):
|
|
wai = os.path.dirname(os.path.abspath(root))
|
|
failures = 0
|
|
nfiles = nvec = 0
|
|
for sub, fn in (("vectors", graph_line), ("stage", stage_line), ("graphs", graph_line)):
|
|
d = os.path.join(root, sub)
|
|
for name in sorted(os.listdir(d)):
|
|
if not name.endswith(".jsonl"):
|
|
continue
|
|
raw = open(os.path.join(d, name), "rb").read()
|
|
out = []
|
|
try:
|
|
for line in raw.decode("ascii").split("\n")[:-1]:
|
|
out.append(cw(fn(json.loads(line))) + "\n")
|
|
except Exception as e:
|
|
failures += 1
|
|
print("FAIL %s/%s: %s: %s" % (sub, name, type(e).__name__, e))
|
|
continue
|
|
got = "".join(out).encode()
|
|
nfiles += 1
|
|
nvec += len(out)
|
|
if write:
|
|
os.makedirs(os.path.join(write, sub), exist_ok=True)
|
|
open(os.path.join(write, sub, name), "wb").write(got)
|
|
if got != raw:
|
|
failures += 1
|
|
a, b = got.split(b"\n"), raw.split(b"\n")
|
|
k = next((k for k in range(min(len(a), len(b))) if a[k] != b[k]), min(len(a), len(b)))
|
|
print("DIFF %s/%s: line %d differs" % (sub, name, k + 1))
|
|
print("rewrote %d vector files, %d vectors" % (nfiles, nvec))
|
|
idx = json.load(open(os.path.join(root, "certificates", "index.json")))
|
|
nsets = ngraphs = 0
|
|
lowered = {}
|
|
for e in idx["sets"]:
|
|
try:
|
|
files = {m: open(os.path.join(wai, p), "rb").read() for m, p in e["members"].items()}
|
|
if set_digest(files) != e["set_sha256"]:
|
|
raise ValueError("the set's digest differs")
|
|
roles = lower_set(e, files)
|
|
if len(roles) != len(e["roles"]):
|
|
raise ValueError("the roles differ")
|
|
for (role, b), path in zip(roles, e["roles"]):
|
|
if not path.endswith("/%s/%s.whg1" % (e["set_sha256"], role)):
|
|
raise ValueError("role %s is at %s" % (role, path))
|
|
if write:
|
|
p = os.path.join(write, os.path.relpath(os.path.join(wai, path), root))
|
|
os.makedirs(os.path.dirname(p), exist_ok=True)
|
|
open(p, "wb").write(b)
|
|
if open(os.path.join(wai, path), "rb").read() != b:
|
|
raise ValueError("%s: the lowered graph differs" % role)
|
|
if e.get("synthetic"):
|
|
lowered[(FORM_FILE[e["lowering"]], e["name"], role)] = b
|
|
ngraphs += 1
|
|
nsets += 1
|
|
except Exception as ex:
|
|
failures += 1
|
|
print("FAIL lowering %s (%s): %s" % (e.get("name"), e.get("capability"), ex))
|
|
print("re-lowered %d sets, %d graphs" % (nsets, ngraphs))
|
|
failures += check_lower_vectors(root, lowered)
|
|
return failures
|
|
|
|
|
|
FORM_FILE = {
|
|
"int_synth/1": "lower_int_synth", "whm1/1": "lower_whm1", "wam1/1": "lower_wam1",
|
|
"wmb1/1": "lower_wmb1", "world/1": "lower_world", "jpegai-int/1": "lower_jpegai",
|
|
}
|
|
|
|
|
|
def check_lower_vectors(root, lowered):
|
|
"""Each vector of graphs/lower_<form>.jsonl has the id <file>/<set>/<role>/...: its
|
|
graph must be that role of that synthetic set as re-lowered here, and every role
|
|
of every synthetic set must have one. `lowered` maps (file, set, role) to bytes."""
|
|
failures = 0
|
|
seen = set()
|
|
d = os.path.join(root, "graphs")
|
|
for name in sorted(os.listdir(d)):
|
|
if not (name.startswith("lower_") and name.endswith(".jsonl")):
|
|
continue
|
|
f = name[:-len(".jsonl")]
|
|
for k, line in enumerate(open(os.path.join(d, name), "rb").read().decode("ascii").split("\n")[:-1]):
|
|
v = json.loads(line)
|
|
parts = v["id"].split("/")
|
|
key = tuple(parts[:3])
|
|
if len(parts) < 4 or parts[0] != f or key not in lowered:
|
|
failures += 1
|
|
print("FAIL graphs/%s line %d: %s names no lowered role" % (name, k + 1, v["id"]))
|
|
elif bytes.fromhex(v["graph"]) != lowered[key]:
|
|
failures += 1
|
|
print("FAIL graphs/%s line %d: %s is not the role as lowered" % (name, k + 1, v["id"]))
|
|
seen.add(key)
|
|
for key in sorted(set(lowered) - seen):
|
|
failures += 1
|
|
print("FAIL %s: no lowering vector for this role" % "/".join(key))
|
|
print("matched %d lowered roles to their lowering vectors" % len(seen & set(lowered)))
|
|
return failures
|
|
|
|
|
|
def main(argv):
|
|
if len(argv) not in (2, 4) or (len(argv) == 4 and argv[2] != "--write"):
|
|
print(__doc__)
|
|
return 2
|
|
write = argv[3] if len(argv) == 4 else None
|
|
failures = check(argv[1], write)
|
|
if failures:
|
|
print("%d failures" % failures)
|
|
return 1
|
|
print("ops-conformance: every vector file and every lowered graph rewritten byte for byte")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv))
|