96f5468eb5
jpegai-int/1 now writes six graphs (hs_y, hs_uv, rec_y, rec_uv, syn_y, syn_uv; integer-operators §13.6.2). The independent checker (tools/wai_ops_check.py), written by a separate session from §13.6.2's text without reading the Rust, lowers the eight weight files and reproduces every lowered graph and certificate byte for byte: the committed synthetic set (4e4ba5e4…) through `ops-conformance` (19 sets, 34 graphs, 19 certificates), and jpegai/v1 (3e243dc1…) through `certify --check`, with `--external jpegai/v1=DIR` recomputing it too. - The verdict: Appendix C's prediction for these networks (acc64) holds; the set meets acc64 and not acc32 (first failure rec_y node 1, conv.acc 42 bits), and the scale graphs still meet acc32 on every node. - §11's role_profiles in both certifiers: `--role-profile ROLE=PROFILE` (exit 1 when the named graph fails it), and the external-set check reads an entry's role_profiles; external.json's jpegai/v1 entry declares acc64 with acc32 for hs_y and hs_uv. wai-jpegai-e2e.yml certifies with both. - The checker's readings of the member rules made two of them explicit: a file with two members of one name at any depth is refused, and a conv member's `w` holds the product of its dims of integers that fit i64. wai-ops' parse and the parsed form's conversion (Weights::ptq) now refuse what the text refuses. - wai-jpegai-e2e.yml gains a step for tests/jpegai_nets_ops.rs on the pinned bundle (and its timeout 60 minutes); both fast-kernel steps of wai-conformance.yml run it on the committed set. - tests/jpegai_ops.rs: the scale-path certificate test reads the set's eleven members and holds the scale graphs, not the set, to acc32. Co-Authored-By: Claude Opus 5.5 <[email protected]>
3253 lines
120 KiB
Python
3253 lines
120 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;
|
|
- is a certifier (section 1): from each set's lowered graphs it computes the
|
|
bounds certificate (the pass of sections 8.1-8.3, the classes and parameter
|
|
conditions of section 7, the flags and reasons of section 9.2, and the canonical
|
|
writer of section 8.5), and requires the bytes of the committed certificate,
|
|
for every set of index.json and the two falsification sets of section 14.3,
|
|
each of which must fail the profile it claims. The sets of external.json are
|
|
certified too when their files are given (--external NAME=DIR); otherwise their
|
|
certificates are checked for what needs no set. Either way, an entry's verdict
|
|
must meet its `profile`, and each graph its `role_profiles` names must meet its
|
|
own profile (section 11).
|
|
|
|
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
|
|
python3 -I tools/wai_ops_check.py ops-conformance --external jpegai/v1=DIR ...
|
|
|
|
Exit status 1 on any difference or failure.
|
|
|
|
python3 -I tools/wai_ops_check.py certify --capability CAP --lowering ID \\
|
|
[--c N] --member NAME=PATH ... [--profile P] [--role-profile ROLE=P ...] \\
|
|
[--out FILE] [--check FILE]
|
|
python3 -I tools/wai_ops_check.py certify --whg1 ROLE=PATH ... [--profile P] ...
|
|
|
|
certifies one set, or WHG1 graphs on their own, writing the certificate to FILE or
|
|
to stdout. --role-profile (repeatable) holds the graph of one role to a profile of
|
|
its own, as a capability's `role_profiles` does (section 11): that graph meets it
|
|
when its inputs' and scalars' declared ranges meet the profile's tensor width and
|
|
every node of it fits the profile (section 8.5). Exit status: 0 the verdict meets
|
|
the profile (acc64 by default) and each named graph meets its own; 1 one does not,
|
|
and the certificate is still written; 2 a graph is malformed or refused at load;
|
|
3 --check finds a difference, and the first differing path is printed with both
|
|
values; 4 the lowering refuses the set; 5 usage, or a --role-profile role that
|
|
names no graph.
|
|
|
|
The lowering jpegai-int/1 (section 13.6) reads eleven members, norm.json and
|
|
weights/{hs_y,hs_uv,hyp_y,hyp_uv,mcm0_y,mcm1_y,mcm2_y,mcm3_y,synth_y,synth_uv}.json,
|
|
and writes six roles: hs_y and hs_uv (13.6.1), rec_y, rec_uv, syn_y and syn_uv
|
|
(13.6.2). Any other member is part of the set's digest and is not read.
|
|
"""
|
|
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 = load_graph(b)
|
|
if not g["verdict"]["acc64"]:
|
|
bad()
|
|
return g
|
|
|
|
|
|
def load_graph(b):
|
|
"""Load steps 1-7, then step 8's bounds pass with the refusals the pass makes
|
|
(lut16's range, an empty `check` interval), but not the profile against the
|
|
verdict: what a certifier needs to write a certificate whatever the verdict."""
|
|
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["pass"] = bounds_pass(g)
|
|
g["verdict"] = verdict_of([("", g)])
|
|
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 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]
|
|
|
|
|
|
PROFILES = ("acc32", "acc64")
|
|
POLICY = ("derived", "clamp", "check")
|
|
_RQ = ["rq.t", "rq.p", "rq.q", "rq.y"]
|
|
# Table 8.5-A: the intermediates each operator records, in definition order. Every
|
|
# other operator records none.
|
|
RECORDS = {
|
|
"conv": ["conv.term", "conv.acc"] + _RQ,
|
|
"conv_transpose": ["convT.term", "convT.acc"] + _RQ,
|
|
"rescale": _RQ,
|
|
"chscale": _RQ,
|
|
"mul": ["mul.p"] + _RQ,
|
|
"gram": ["gram.term", "gram.acc"] + _RQ,
|
|
"apply": ["apply.term", "apply.acc"] + _RQ,
|
|
"leaky": ["leaky.p"],
|
|
"igdn": ["igdn.x2", "igdn.gx2", "igdn.term", "igdn.N", "igdn.r", "igdn.xr"],
|
|
"layernorm_c": ["ln.S", "ln.m", "ln.d", "ln.T", "ln.v", "rsq.num", "ln.k", "ln.dk", "ln.n", "ln.ng"],
|
|
"l2norm": ["l2.T", "rsq.num", "l2.k", "l2.vk"],
|
|
"softmax": ["sm.diff", "exp.negx", "exp.nl", "exp.t2", "exp.c", "sm.sum", "sm.num"],
|
|
"sigmoid": ["exp.negx", "exp.nl", "exp.t2", "exp.c", "sig.num"],
|
|
"elu": ["exp.negx", "exp.nl", "exp.t2", "exp.c", "elu.p"],
|
|
}
|
|
# Section 8.5, k_c: the operators that record it, and the term interval of section 6
|
|
# it is taken over.
|
|
KC_TERM = {"conv": "conv.term", "conv_transpose": "convT.term", "igdn": "igdn.term", "gram": "gram.term", "apply": "apply.term"}
|
|
# Section 8.5, mac: the operators that record it; w_bits for the first three only.
|
|
MAC = ("conv", "conv_transpose", "chscale", "mul", "gram", "apply")
|
|
# Section 5's table: the operators marked "no" in the wai-ops/1-npu column.
|
|
NPU_NO = ("igdn", "layernorm_c", "l2norm", "softmax", "sigmoid", "elu", "gram", "apply")
|
|
|
|
|
|
def bits(iv):
|
|
"""bits([lo, hi]) of section 2: the smallest N >= 1 such that both fit iN."""
|
|
return max((v if v >= 0 else ~v).bit_length() + 1 for v in iv)
|
|
|
|
|
|
def acc32_params(op):
|
|
"""The parameter conditions of section 7.2 the node fails, as their keys, in that
|
|
table's order."""
|
|
n, P = op["name"], op["p"]
|
|
f32 = lambda v: I32_MIN <= v <= I32_MAX
|
|
keys = []
|
|
rq = P.get("rq")
|
|
if rq is not None:
|
|
if not all(map(f32, rq["mult"])):
|
|
keys.append("rq.mult")
|
|
if not all(s <= 63 for s in rq["shift"]):
|
|
keys.append("rq.shift")
|
|
if rq["bias"] == 2 and not all(map(f32, rq["biasv"])):
|
|
keys.append("rq.post_bias")
|
|
if n in ("conv", "conv_transpose") and not all(map(f32, P["w"])):
|
|
keys.append("weights")
|
|
if n == "chscale":
|
|
if not all(map(f32, P["M"])):
|
|
keys.append("chscale.m")
|
|
if P["shift"] > 63:
|
|
keys.append("chscale.s")
|
|
if n == "igdn":
|
|
if not all(map(f32, P["gamma"])):
|
|
keys.append("igdn.gamma")
|
|
if P["gf"] > 63 or P["frac"] > 63:
|
|
keys.append("igdn.shift")
|
|
if n == "leaky" and not f32(P["a"]):
|
|
keys.append("leaky.a")
|
|
return keys
|
|
|
|
|
|
def node_npu(g, op):
|
|
"""The node's flags of section 9.2, in ascending order, and its reasons of section
|
|
8.5 other than "acc32"."""
|
|
n, P = op["name"], op["p"]
|
|
flags, reasons = set(), set()
|
|
if n == "lut16":
|
|
flags.add("lut16")
|
|
if n in ("add", "chunk_add"):
|
|
state = {i["tensor"] for i in g["inputs"] if i["kind"] == 1}
|
|
if any(t in state for t in op["ins"]):
|
|
flags.add("state_add")
|
|
rq = None
|
|
if n == "chscale":
|
|
flags.add("chscale")
|
|
# A chscale always has a multiplier, M[r][c], and no bias.
|
|
rq = (P["round"], True, P["M"], [P["shift"]], False)
|
|
elif "rq" in P:
|
|
q = P["rq"]
|
|
rq = (q["round"], q["m_kind"] != 0, q["mult"], q["shift"], q["bias"] == 2)
|
|
if rq is not None:
|
|
rnd, has_mult, mults, shifts, post = rq
|
|
if rnd == 2:
|
|
flags.add("rq_half_even")
|
|
if rnd == 0 and has_mult:
|
|
flags.add("rq_floor_mult")
|
|
if post:
|
|
flags.add("rq_post_bias")
|
|
if any(m < 0 for m in mults):
|
|
flags.add("rq_neg_mult")
|
|
if any(not 0 <= s <= 31 for s in shifts):
|
|
flags.add("rq_shift_range")
|
|
if n in NPU_NO:
|
|
reasons.add("op:" + n)
|
|
if n == "conv":
|
|
cin, G = g["tensors"][op["ins"][0]]["c"], P["groups"]
|
|
if (G != 1 and G != cin) or (G == cin and G > 1 and P["cout"] != cin):
|
|
reasons.add("conv_groups")
|
|
if n == "crop" and any(o[0] == "scalar" for o in P["ops"]):
|
|
reasons.add("crop_scalar")
|
|
return sorted(flags), reasons
|
|
|
|
|
|
def width_failures(nd, p):
|
|
"""The node's recorded intermediates and producer results (`out.<k>`) that do not
|
|
meet profile p's width for their class (section 7.1), as (bits, key)."""
|
|
W = WIDTH[p]
|
|
f = [(bits(r), k) for k, r in nd["rec"].items() if bits(r) > W[CLASS[k]]]
|
|
f += [(bits(o["result"]), "out.%d" % j) for j, o in enumerate(nd["outs"]) if bits(o["result"]) > W["tensor"]]
|
|
return f
|
|
|
|
|
|
def bounds_pass(g):
|
|
"""The bounds pass of sections 8.1-8.3 over a graph that passed load steps 1-7,
|
|
in graph order. Refuses (model_mismatch) where the pass refuses the graph at
|
|
load: a `check` interval that is empty (section 8.1), or lut16's input range
|
|
(section 5.6). Returns one record per node: its recorded intermediates, whether
|
|
it saturated, its input tensors' intervals, each output's producer result and
|
|
interval after its policy, its acc32_params, its npu flags and reasons, and the
|
|
profiles it fits (section 8.5)."""
|
|
T = g["tensors"]
|
|
iv = {i["tensor"]: (T[i["tensor"]]["lo"], T[i["tensor"]]["hi"]) for i in g["inputs"]}
|
|
nodes = []
|
|
for k, op in enumerate(g["ops"]):
|
|
nb = Node()
|
|
ins = [iv[t] for t in op["ins"]]
|
|
results = node_bounds(nb, g, op, iv)
|
|
if sorted(nb.rec) != sorted(RECORDS.get(op["name"], [])):
|
|
raise AssertionError("%s records %s" % (op["name"], sorted(nb.rec)))
|
|
outs = []
|
|
for t, r in zip(op["outs"], results):
|
|
rec = T[t]
|
|
if rec["policy"] == 0:
|
|
ti = r
|
|
elif rec["policy"] == 1:
|
|
ti = (min(max(r[0], rec["lo"]), rec["hi"]), min(max(r[1], rec["lo"]), rec["hi"]))
|
|
else:
|
|
ti = (max(r[0], rec["lo"]), min(r[1], rec["hi"]))
|
|
if ti[0] > ti[1]:
|
|
bad()
|
|
iv[t] = ti
|
|
outs.append({"tensor": t, "iv": ti, "result": r, "policy": rec["policy"]})
|
|
flags, reasons = node_npu(g, op)
|
|
nd = {"index": k, "op": op, "rec": nb.rec, "saturated": nb.saturated, "ins": ins, "outs": outs,
|
|
"acc32_params": acc32_params(op), "flags": flags, "reasons": reasons}
|
|
nd["fits"] = {}
|
|
for p in PROFILES:
|
|
fit = not nd["saturated"] and not width_failures(nd, p)
|
|
if p == "acc32":
|
|
fit = fit and not nd["acc32_params"] and all(o["policy"] != 2 for o in outs)
|
|
nd["fits"][p] = fit
|
|
nodes.append(nd)
|
|
return nodes
|
|
|
|
|
|
def node_failure(nd, p, name):
|
|
"""The failure of section 8.5 for a node that does not fit profile p."""
|
|
f = {"graph": name, "index": nd["index"], "where": "node"}
|
|
if nd["saturated"]:
|
|
f["key"] = "saturated"
|
|
return f
|
|
w = width_failures(nd, p)
|
|
if w:
|
|
f["bits"], f["key"] = min(w, key=lambda e: (-e[0], e[1].encode()))
|
|
else:
|
|
f["key"] = nd["acc32_params"][0] if nd["acc32_params"] else "policy.check"
|
|
return f
|
|
|
|
|
|
def verdict_of(graphs):
|
|
"""The verdict of section 8.5 over [(name, g)], each g carrying its pass."""
|
|
v = {}
|
|
for p in PROFILES:
|
|
tw = WIDTH[p]["tensor"]
|
|
fail = None
|
|
for name, g in graphs:
|
|
T = g["tensors"]
|
|
for k, i in enumerate(g["inputs"]):
|
|
r = (T[i["tensor"]]["lo"], T[i["tensor"]]["hi"])
|
|
if bits(r) > tw:
|
|
fail = {"bits": bits(r), "graph": name, "index": k, "key": "range", "where": "input"}
|
|
break
|
|
if fail is None:
|
|
for k, r in enumerate(g["scalars"]):
|
|
if bits(r) > tw:
|
|
fail = {"bits": bits(r), "graph": name, "index": k, "key": "range", "where": "scalar"}
|
|
break
|
|
if fail is None:
|
|
for nd in g["pass"]:
|
|
if not nd["fits"][p]:
|
|
fail = node_failure(nd, p, name)
|
|
break
|
|
if fail is not None:
|
|
break
|
|
v[p] = fail is None
|
|
if fail is not None:
|
|
v[p + "_first_failure"] = fail
|
|
reasons, flags = set(), set()
|
|
if not v["acc32"]:
|
|
reasons.add("acc32")
|
|
for _, g in graphs:
|
|
for nd in g["pass"]:
|
|
reasons |= nd["reasons"]
|
|
flags |= set(nd["flags"])
|
|
v["npu_reasons"] = sorted(reasons, key=lambda s: s.encode())
|
|
v["npu"] = not reasons
|
|
v["npu_flags"] = sorted(flags, key=lambda s: s.encode())
|
|
return v
|
|
|
|
|
|
# --------------------------------------------------------------- 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; should one be formed
|
|
# all the same, the decode is refused rather than carried on.
|
|
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
|
|
|
|
|
|
# --------------------------------------------------------------- the certificate (8.5)
|
|
|
|
|
|
def ivj(iv):
|
|
return {"bits": bits(iv), "hi": iv[1], "lo": iv[0]}
|
|
|
|
|
|
def node_json(nd):
|
|
op = nd["op"]
|
|
n, P = op["name"], op["p"]
|
|
d = {
|
|
"acc32_params": nd["acc32_params"],
|
|
"fits": [p for p in PROFILES if nd["fits"][p]],
|
|
"index": nd["index"],
|
|
"intermediates": {k: ivj(nd["rec"][k]) for k in RECORDS.get(n, [])},
|
|
"npu_flags": nd["flags"],
|
|
"op": n,
|
|
"out": [{"bits": bits(o["iv"]), "hi": o["iv"][1], "lo": o["iv"][0], "policy": POLICY[o["policy"]],
|
|
"result": ivj(o["result"]), "tensor": o["tensor"]} for o in nd["outs"]],
|
|
"saturated": nd["saturated"],
|
|
}
|
|
if n in KC_TERM:
|
|
t = nd["rec"][KC_TERM[n]]
|
|
m = max(abs(t[0]), abs(t[1]))
|
|
d["k_c"] = "unbounded" if m == 0 else I32_MAX // m
|
|
if n in MAC:
|
|
d["mac"] = {"a_bits": [bits(x) for x in nd["ins"]]}
|
|
if n in ("conv", "conv_transpose", "chscale"):
|
|
w = P["M"] if n == "chscale" else P["w"]
|
|
d["mac"]["w_bits"] = bits((min(w), max(w)))
|
|
return d
|
|
|
|
|
|
def graph_json(name, b, g):
|
|
T = g["tensors"]
|
|
inputs = []
|
|
for i in g["inputs"]:
|
|
rec = T[i["tensor"]]
|
|
e = {"hi": rec["hi"], "kind": "state" if i["kind"] else "data", "lo": rec["lo"], "max_h": i["max_h"],
|
|
"max_w": i["max_w"], "policy": POLICY[rec["policy"]], "tensor": i["tensor"]}
|
|
if i["kind"]:
|
|
e["init"], e["src"] = i["init"], i["src"]
|
|
inputs.append(e)
|
|
return {
|
|
"graph_sha256": hashlib.sha256(b).hexdigest(),
|
|
"inputs": inputs,
|
|
"name": name,
|
|
"nodes": [node_json(nd) for nd in g["pass"]],
|
|
"scalars": [{"hi": hi, "index": k, "lo": lo} for k, (lo, hi) in enumerate(g["scalars"])],
|
|
}
|
|
|
|
|
|
def canonical(v):
|
|
"""The canonical JSON of section 8.5, as bytes. Every string is printable ASCII;
|
|
a role or capability name outside that alphabet cannot be certified."""
|
|
def strings(x):
|
|
if isinstance(x, str):
|
|
yield x
|
|
elif isinstance(x, list):
|
|
for y in x:
|
|
yield from strings(y)
|
|
elif isinstance(x, dict):
|
|
for k, y in x.items():
|
|
yield k
|
|
yield from strings(y)
|
|
for s in strings(v):
|
|
if any(not 0x20 <= ord(ch) <= 0x7E for ch in s):
|
|
raise ValueError("a string outside printable ASCII cannot be certified: %r" % s)
|
|
return cw(v).encode("ascii")
|
|
|
|
|
|
def certificate(capability, lowering, set_sha256, roles, role_profiles=None):
|
|
"""The bounds certificate of section 8.5 for the graphs [(role, WHG1 bytes)], in
|
|
the order given: (its canonical bytes, its verdict). Each graph passes load steps
|
|
1-7 and the bounds pass, or the Refuse of the step that fails propagates. With
|
|
`role_profiles` ({role: profile}, section 11), it returns a third value, the
|
|
failure of section 8.5 of each named graph that does not meet its profile, judged
|
|
on that graph alone ({role: failure}; empty when every one meets it)."""
|
|
gs = [(name, b, load_graph(b)) for name, b in roles]
|
|
v = verdict_of([(name, g) for name, _, g in gs])
|
|
cert = {
|
|
"capability": capability,
|
|
"certificate": "wai-ops/1 bounds",
|
|
"graphs": [graph_json(name, b, g) for name, b, g in gs],
|
|
"lowering": lowering,
|
|
"set_sha256": set_sha256,
|
|
"verdict": v,
|
|
}
|
|
if role_profiles is None:
|
|
return canonical(cert), v
|
|
fails = {}
|
|
for name, p in role_profiles.items():
|
|
# A graph meets a profile when its inputs' and scalars' declared ranges meet
|
|
# the profile's tensor width and every node of it fits the profile: the
|
|
# verdict of section 8.5 over that graph alone.
|
|
rv = verdict_of([(n, g) for n, _, g in gs if n == name])
|
|
if not rv[p]:
|
|
fails[name] = rv[p + "_first_failure"]
|
|
return canonical(cert), v, fails
|
|
|
|
|
|
ABSENT = "(absent)"
|
|
|
|
|
|
def json_diffs(a, b, path="", out=None, limit=None):
|
|
"""The JSON paths at which two parsed values differ, in document order, each
|
|
with both values: [(path, a's value, b's value)]."""
|
|
out = [] if out is None else out
|
|
if limit is not None and len(out) >= limit:
|
|
return out
|
|
if isinstance(a, dict) and isinstance(b, dict):
|
|
for k in sorted(set(a) | set(b), key=lambda s: s.encode()):
|
|
p = path + ("." + k if k.isidentifier() else '["%s"]' % k)
|
|
if k not in a or k not in b:
|
|
out.append((p, a.get(k, ABSENT), b.get(k, ABSENT)))
|
|
else:
|
|
json_diffs(a[k], b[k], p, out, limit)
|
|
elif isinstance(a, list) and isinstance(b, list):
|
|
for k in range(max(len(a), len(b))):
|
|
p = "%s[%d]" % (path, k)
|
|
if k >= len(a) or k >= len(b):
|
|
out.append((p, a[k] if k < len(a) else ABSENT, b[k] if k < len(b) else ABSENT))
|
|
else:
|
|
json_diffs(a[k], b[k], p, out, limit)
|
|
elif a != b:
|
|
out.append((path or "$", a, b))
|
|
return out[:limit] if limit is not None else out
|
|
|
|
|
|
def first_difference(got, want):
|
|
"""Part of what a `--check` mismatch prints: the first differing JSON path, with
|
|
both values."""
|
|
try:
|
|
a, b = json.loads(got), json.loads(want)
|
|
except ValueError:
|
|
return "the committed file is not JSON"
|
|
d = json_diffs(a, b, limit=1)
|
|
if not d:
|
|
return "the same JSON, written differently"
|
|
p, x, y = d[0]
|
|
return "first difference at %s: computed %s, committed %s" % (p, cw(x) if x is not ABSENT else x, cw(y) if y is not ABSENT else y)
|
|
|
|
|
|
# --------------------------------------------------------------- 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 op_n(self, code, ins, out_recs, params):
|
|
"""An operator of several outputs, which take the next indices in their order."""
|
|
ts = [self.tensor(*r) for r in out_recs]
|
|
self.ops.append((code, ins, ts, params))
|
|
return ts
|
|
|
|
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))]
|
|
|
|
|
|
JPEGAI_MEMBERS = (
|
|
"norm.json", "weights/hs_y.json", "weights/hs_uv.json", "weights/hyp_y.json", "weights/hyp_uv.json",
|
|
"weights/mcm0_y.json", "weights/mcm1_y.json", "weights/mcm2_y.json", "weights/mcm3_y.json",
|
|
"weights/synth_y.json", "weights/synth_uv.json",
|
|
)
|
|
|
|
|
|
def jpegai_lower(files):
|
|
"""jpegai-int/1 (section 13.6): the scale graphs of 13.6.1, then the graphs of
|
|
13.6.2, in section 13.0's role order. `files` maps member names to bytes; a member
|
|
the lowering reads that the set lacks refuses it. Other members are not read."""
|
|
if any(m not in files for m in JPEGAI_MEMBERS):
|
|
lref()
|
|
roles = jpegai_scale_lower(files["norm.json"], files["weights/hs_y.json"], files["weights/hs_uv.json"])
|
|
W = {k: q16_file(files["weights/%s.json" % k]) for k in ("hyp_y", "hyp_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "synth_y", "synth_uv")}
|
|
roles.append(("rec_y", rec_graph(W["hyp_y"], [W["mcm%d_y" % i] for i in range(4)], 160)))
|
|
roles.append(("rec_uv", rec_graph(W["hyp_uv"], None, 96)))
|
|
roles.append(("syn_y", syn_y_graph(W["synth_y"])))
|
|
roles.append(("syn_uv", syn_uv_graph(W["synth_uv"])))
|
|
return roles
|
|
|
|
|
|
def q16_file(b):
|
|
"""One weights file of section 13.6.2, *The members*: {name: member} for each
|
|
member whose value is an object of kind "conv" or "q16"; every other member is
|
|
ignored. Refused: a file that is not a JSON object, two members of one object with
|
|
one name (anywhere in the file), and a member of either kind that lacks one of its
|
|
fields or holds another value."""
|
|
def pairs(ps):
|
|
if len({k for k, _ in ps}) != len(ps):
|
|
lref()
|
|
return dict(ps)
|
|
|
|
try:
|
|
root = json.loads(b, object_pairs_hook=pairs)
|
|
except ValueError:
|
|
lref()
|
|
if not isinstance(root, dict):
|
|
lref()
|
|
isint = lambda v: isinstance(v, int) and not isinstance(v, bool)
|
|
out = {}
|
|
for name, m in root.items():
|
|
if not isinstance(m, dict) or m.get("kind") not in ("conv", "q16"):
|
|
continue
|
|
if m["kind"] == "conv":
|
|
dims, w, mult, shift = m.get("dims"), m.get("w"), m.get("mult"), m.get("shift")
|
|
if not isinstance(dims, list) or not all(isint(d) and d >= 0 for d in dims):
|
|
lref()
|
|
n = 1
|
|
for d in dims:
|
|
n *= d
|
|
if not isinstance(w, list) or len(w) != n or not all(isint(x) and I64_MIN <= x <= I64_MAX for x in w):
|
|
lref()
|
|
if not isint(mult) or not I64_MIN <= mult <= I64_MAX or not isint(shift) or not 0 <= shift <= 127:
|
|
lref()
|
|
out[name] = {"kind": "conv", "dims": dims, "w": w, "mult": mult, "shift": shift}
|
|
else:
|
|
v = m.get("v")
|
|
if not isinstance(v, list) or not all(isint(x) and I64_MIN <= x <= I64_MAX for x in v):
|
|
lref()
|
|
out[name] = {"kind": "q16", "v": v}
|
|
return out
|
|
|
|
|
|
A32 = (2, I32_MIN, I32_MAX)
|
|
Q16_RQ = rq_block(0, 0, [], 0, 0, [16], 1, [], 0) # mul: no multiplier, shift 16, no bias, floor
|
|
ATT_RQ = rq_block(1, 0, [], 0, 0, [16], 1, [], 0) # gram, apply: shift 16, half_up
|
|
SUBLATTICE = bytes([0, 3, 1, 2]) # (0,0), (1,1), (0,1), (1,0), each 2a + b
|
|
|
|
|
|
class Q16Net:
|
|
"""The blocks of section 13.6.2 over one graph under construction, reading the
|
|
modules of the weights file F (which a graph may switch, as rec_y does)."""
|
|
|
|
def __init__(self):
|
|
self.wb = WB()
|
|
self.F = {}
|
|
|
|
def c(self, x):
|
|
return self.wb.tensors[x][0]
|
|
|
|
def weight(self, m):
|
|
e = self.F.get(m + ".weight")
|
|
if e is None or e["kind"] != "conv":
|
|
lref()
|
|
return e
|
|
|
|
def q16(self, name):
|
|
e = self.F.get(name)
|
|
if e is None or e["kind"] != "q16":
|
|
lref()
|
|
return e["v"]
|
|
|
|
def bias(self, m, cout):
|
|
"""m.bias, cout values, or None when the file has none."""
|
|
e = self.F.get(m + ".bias")
|
|
if e is None:
|
|
return None
|
|
if e["kind"] != "q16" or len(e["v"]) != cout:
|
|
lref()
|
|
return e["v"]
|
|
|
|
def rq(self, e, shift, bias):
|
|
"""C's RQ: floor; one multiplier; one shift; a post bias, or none."""
|
|
mult = e["mult"]
|
|
if bias is None:
|
|
return rq_block(0, 0, [mult], 1, smallest([mult], (4, 8)), [shift], 1, [], 0)
|
|
return rq_block(0, 2, [mult], 1, smallest([mult], (4, 8)), [shift], 1, bias, smallest(bias, (4, 8)))
|
|
|
|
def wbytes(self, w):
|
|
if any(not I64_MIN <= v <= I64_MAX for v in w):
|
|
lref()
|
|
return smallest(w, (1, 2, 4, 8))
|
|
|
|
def conv(self, x, m, k, s, p, g, shift=None):
|
|
"""C(m; k, s, p, g); g = "dw" is groups = c, with cout = c."""
|
|
c = self.c(x)
|
|
e = self.weight(m)
|
|
dims = e["dims"]
|
|
dw = g == "dw"
|
|
if dw:
|
|
g = c
|
|
if len(dims) != 4 or c % g:
|
|
lref()
|
|
cout = dims[0]
|
|
if cout == 0 or cout % g or dims != [cout, c // g, k, k] or (dw and cout != c):
|
|
lref()
|
|
rq = self.rq(e, e["shift"] if shift is None else shift, self.bias(m, cout))
|
|
w = e["w"]
|
|
return self.wb.op(0x01, [x], (cout,) + A32, p_conv(cout, k, k, s, s, p, p, g, self.wbytes(w), rq, w))
|
|
|
|
def convt(self, x, m, k, op_):
|
|
"""T(m; k, op): stride 2, padding 1, output padding op, with C's RQ."""
|
|
c = self.c(x)
|
|
e = self.weight(m)
|
|
dims = e["dims"]
|
|
if len(dims) != 4 or dims[0] != c or dims[1] == 0 or dims[2:] != [k, k]:
|
|
lref()
|
|
cout = dims[1]
|
|
rq = self.rq(e, e["shift"], self.bias(m, cout))
|
|
w = e["w"]
|
|
return self.wb.op(0x11, [x], (cout,) + A32, p_convt(cout, k, k, 2, 2, 1, 1, op_, op_, self.wbytes(w), rq, w))
|
|
|
|
def same(self, *xs):
|
|
"""The inputs of each mul and each add have one channel count."""
|
|
if len({self.c(x) for x in xs}) != 1:
|
|
lref()
|
|
return self.c(xs[0])
|
|
|
|
def add(self, a, b):
|
|
return self.wb.op(0x05, [a, b], (self.same(a, b),) + A32, b"")
|
|
|
|
def mul(self, a, b):
|
|
return self.wb.op(0x10, [a, b], (self.same(a, b),) + A32, Q16_RQ)
|
|
|
|
def unary(self, code, x, params=b""):
|
|
return self.wb.op(code, [x], (self.c(x),), params)
|
|
|
|
def relu6(self, x):
|
|
return self.unary(0x0B, x, p_clamp(0, 6 << 16))
|
|
|
|
def split(self, x, start, count):
|
|
if start + count > self.c(x):
|
|
lref()
|
|
return self.wb.op(0x08, [x], (count,), struct.pack("<II", start, count))
|
|
|
|
def concat(self, xs):
|
|
return self.wb.op(0x06, list(xs), (sum(self.c(x) for x in xs),), b"")
|
|
|
|
def shuffle2(self, x):
|
|
if self.c(x) % 4:
|
|
lref()
|
|
return self.wb.op(0x09, [x], (self.c(x) // 4,), struct.pack("<I", 2))
|
|
|
|
def merge(self, xs):
|
|
return self.wb.op(0x16, list(xs), (self.same(*xs),), SUBLATTICE)
|
|
|
|
def R(self, x, m):
|
|
t = self.conv(x, m + ".conv1", 3, 1, 1, 1)
|
|
t = self.unary(0x0C, t)
|
|
t = self.conv(t, m + ".conv2", 3, 1, 1, 1)
|
|
return self.add(x, t)
|
|
|
|
def G(self, x, m):
|
|
"""The gate x·(f(x) + 1)."""
|
|
t = self.relu6(x)
|
|
t = self.conv(t, m + ".conv", 3, 1, 1, 4)
|
|
t = self.conv(t, m + ".conv2", 1, 1, 0, 1)
|
|
c = self.c(t)
|
|
t = self.wb.op(0x02, [t], (c,) + A32, rq_block(0, 2, [], 0, 0, [0], 1, [1 << 16] * c, 4))
|
|
return self.mul(x, t)
|
|
|
|
def CAB(self, x):
|
|
t = self.R(x, "CAB.residual_trunk.0")
|
|
t = self.R(t, "CAB.residual_trunk.1")
|
|
s = self.conv(x, "CAB.subscale", 3, 2, 1, 1)
|
|
s = self.R(s, "CAB.residual_mask2.0")
|
|
s = self.R(s, "CAB.residual_mask3.0")
|
|
s = self.convt(s, "CAB.upscale", 3, 1)
|
|
s = self.unary(0x1B, s)
|
|
return self.add(x, self.mul(t, s))
|
|
|
|
def D(self, x):
|
|
"""The denormalisation, clamp [0, 255·2^16]."""
|
|
c = self.c(x)
|
|
rq = rq_block(0, 2, [128], 1, 4, [0], 1, [127 * (1 << 16) + (1 << 15)] * c, 4)
|
|
return self.wb.op(0x02, [x], (c, 1, 0, 255 << 16), rq)
|
|
|
|
def layernorm(self, x, m):
|
|
c = self.c(x)
|
|
gm, bt = self.q16(m + ".weight"), self.q16(m + ".bias")
|
|
if len(gm) != c or len(bt) != c:
|
|
lref()
|
|
gb = smallest(gm, (4, 8))
|
|
return self.wb.op(0x18, [x], (c,) + A32, struct.pack("<QB7x", 42949, gb) + ints_le(gm, gb) + ints_le(bt, 8))
|
|
|
|
def TAB(self, x, m):
|
|
c = self.c(x)
|
|
if c % 4:
|
|
lref()
|
|
d = c // 4
|
|
sa = self.weight(m + ".prep_data.conv1")["shift"]
|
|
sf = self.weight(m + ".ffn.project_in")["shift"]
|
|
t = self.layernorm(x, m + ".prep_data.norm1")
|
|
t = self.conv(t, m + ".prep_data.conv1", 1, 1, 0, 1, sa)
|
|
if self.c(t) != 3 * c:
|
|
lref()
|
|
t = self.conv(t, m + ".prep_data.conv2", 3, 1, 1, "dw", sa)
|
|
q, k, v = self.split(t, 0, c), self.split(t, c, c), self.split(t, 2 * c, c)
|
|
q = self.unary(0x19, q, struct.pack("<Q", 42949))
|
|
k = self.unary(0x19, k, struct.pack("<Q", 42949))
|
|
a = self.wb.op(0x1D, [q, k], (c,) + A32, struct.pack("<I", 4) + ATT_RQ)
|
|
temp = self.q16(m + ".attn.temperature")
|
|
if len(temp) != 4:
|
|
lref()
|
|
mult = [temp[i // d] for i in range(c)]
|
|
a = self.wb.op(0x02, [a], (c,) + A32, rq_block(1, 0, mult, 2, smallest(mult, (4, 8)), [16], 1, [], 0))
|
|
a = self.unary(0x1A, a)
|
|
a = self.wb.op(0x1E, [a, v], (c,) + A32, struct.pack("<I", 4) + ATT_RQ)
|
|
a = self.conv(a, m + ".attn.project_out", 1, 1, 0, 1, sa)
|
|
x2 = self.add(x, a)
|
|
t = self.layernorm(x2, m + ".ffn.norm1")
|
|
t = self.conv(t, m + ".ffn.project_in", 1, 1, 0, 1, sf)
|
|
if self.c(t) % 2:
|
|
lref()
|
|
h = self.c(t) // 2
|
|
t = self.conv(t, m + ".ffn.dwconv", 3, 1, 1, "dw", sf)
|
|
a, b = self.split(t, 0, h), self.split(t, h, h)
|
|
u = self.mul(self.unary(0x1C, a, struct.pack("<q", 1 << 16)), b)
|
|
u = self.conv(u, m + ".ffn.project_out", 1, 1, 0, 1, sf)
|
|
return self.add(x2, u)
|
|
|
|
|
|
def rec_graph(hyp, mcms, C):
|
|
"""rec_y (mcms: the four context stages' files) or rec_uv (mcms None), 13.6.2."""
|
|
n = Q16Net()
|
|
n.F = hyp
|
|
z = n.wb.input(C, -31, 32, 4, 4)
|
|
r = n.wb.input(C, I32_MIN, I32_MAX, 16, 16)
|
|
t = n.wb.op(0x02, [z], (C,), rq_block(0, 0, [1 << 16], 1, 4, [0], 1, [], 0))
|
|
t = n.conv(t, "conv1", 1, 1, 0, 1)
|
|
t = n.convt(t, "conv2", 4, 0)
|
|
t = n.relu6(t)
|
|
t = n.conv(t, "conv3", 3, 1, 1, 1)
|
|
t = n.relu6(t)
|
|
psi = n.conv(t, "conv4", 3, 1, 1, 1)
|
|
if n.c(psi) != 4 * C:
|
|
lref()
|
|
if mcms is None:
|
|
out = n.add(r, n.merge([n.split(psi, i * C, C) for i in range(4)]))
|
|
return n.wb.bytes([out])
|
|
rs = n.wb.op_n(0x15, [r], [(C,)] * 4, SUBLATTICE)
|
|
ys = []
|
|
for i in range(4):
|
|
n.F = mcms[i]
|
|
p = n.split(psi, i * C, C)
|
|
if i == 0:
|
|
f = p
|
|
else:
|
|
ctx = ys[0] if i == 1 else n.concat(ys)
|
|
f = n.conv(ctx, "conv.0", 1, 1, 0, 1)
|
|
f = n.conv(f, "conv.1", 3, 1, 1, 5)
|
|
f = n.concat([f, p])
|
|
f = n.unary(0x0C, n.conv(f, "fusion_pred_net.conv1", 1, 1, 0, 1))
|
|
f = n.unary(0x0C, n.conv(f, "fusion_pred_net.conv2", 1, 1, 0, 1))
|
|
mu = n.conv(f, "fusion_pred_net.conv3", 1, 1, 0, 1)
|
|
if n.c(mu) != C:
|
|
lref()
|
|
ys.append(n.add(rs[i], mu))
|
|
return n.wb.bytes([n.merge(ys)])
|
|
|
|
|
|
def syn_y_graph(F):
|
|
n = Q16Net()
|
|
n.F = F
|
|
x = n.wb.input(160, I32_MIN, I32_MAX, 16, 16)
|
|
x = n.R(x, "first_stage.0.0")
|
|
x = n.convt(x, "first_stage.1", 3, 1)
|
|
x = n.G(x, "first_stage.2")
|
|
x = n.convt(x, "conv2_t", 3, 1)
|
|
x = n.CAB(x)
|
|
x = n.G(x, "iact2")
|
|
x = n.conv(x, "conv3_t", 1, 1, 0, 1)
|
|
x = n.shuffle2(x)
|
|
x = n.conv(x, "TAM.ds_conv", 3, 2, 1, 1)
|
|
x = n.TAB(x, "TAM.TABs.0")
|
|
x = n.TAB(x, "TAM.TABs.1")
|
|
x = n.convt(x, "TAM.us_conv", 3, 1)
|
|
x = n.G(x, "iact3")
|
|
x = n.convt(x, "conv4_t", 3, 1)
|
|
if n.c(x) != 1:
|
|
lref()
|
|
return n.wb.bytes([n.D(x)])
|
|
|
|
|
|
def syn_uv_graph(F):
|
|
n = Q16Net()
|
|
n.F = F
|
|
y = n.wb.input(160, I32_MIN, I32_MAX, 16, 16)
|
|
uv = n.wb.input(96, I32_MIN, I32_MAX, 16, 16)
|
|
x = n.concat([y, uv])
|
|
x = n.conv(x, "first_stage.conv1", 3, 1, 1, 1)
|
|
x = n.concat([x, uv])
|
|
x = n.convt(x, "conv2_t", 3, 1)
|
|
x = n.CAB(x)
|
|
x = n.G(x, "iact2")
|
|
x = n.conv(x, "conv3_t", 1, 1, 0, 1)
|
|
x = n.shuffle2(x)
|
|
x = n.TAB(x, "TAM.TABs.0")
|
|
x = n.TAB(x, "TAM.TABs.1")
|
|
x = n.G(x, "iact3")
|
|
x = n.convt(x, "conv4_t", 3, 1)
|
|
if n.c(x) != 8:
|
|
lref()
|
|
x = n.shuffle2(x)
|
|
return n.wb.bytes([n.D(x)])
|
|
|
|
|
|
def jpegai_scale_lower(norm_b, hy_b, huv_b):
|
|
"""jpegai-int/1, section 13.6.1: the scale graphs hs_y and hs_uv."""
|
|
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):
|
|
"""The set's lowered graphs, [(role, WHG1 bytes)] in section 13.0's role order.
|
|
Raises LowerRefuse when the lowering refuses the set."""
|
|
low, cap = entry["lowering"], entry["capability"]
|
|
if low == "int_synth/1":
|
|
return int_synth_lower(files["int_synth.json"])
|
|
if low == "whm1/1" and cap == "wai.splat.int_codec":
|
|
return [("g_s", whm1_lower(files["g_s.whm"], True))]
|
|
if low == "whm1/1":
|
|
return [("h_s", whm1_lower(files["h_s.whm"], False)), ("g_s", whm1_lower(files["g_s.whm"], True))]
|
|
if low == "wam1/1":
|
|
return [("g_s", wam1_lower(files["model.wam"]))]
|
|
if low == "wmb1/1":
|
|
return wmb1_lower(files["model.wmb"])
|
|
if low == "world/1":
|
|
return world_lower(files["T.whm"], files["g_s.whm"], dec(entry["c"]))
|
|
if low == "jpegai-int/1":
|
|
return jpegai_lower(files)
|
|
raise ValueError("no lowering %r" % low)
|
|
|
|
|
|
def cert_path(root, entry):
|
|
"""Section 14.3: certificates/<capability>/<set digest>.json."""
|
|
return os.path.join(root, "certificates", entry["capability"], entry["set_sha256"] + ".json")
|
|
|
|
|
|
def check_set_certificate(root, entry, roles, write):
|
|
"""Recomputes a set's certificate from its lowered graphs and requires the bytes
|
|
of the committed one. Returns the verdict."""
|
|
got, v = certificate(entry["capability"], entry["lowering"], entry["set_sha256"], roles)
|
|
path = cert_path(root, entry)
|
|
if write:
|
|
p = os.path.join(write, os.path.relpath(path, root))
|
|
os.makedirs(os.path.dirname(p), exist_ok=True)
|
|
open(p, "wb").write(got)
|
|
want = open(path, "rb").read()
|
|
if got != want:
|
|
raise ValueError("the certificate differs: %s" % first_difference(got, want))
|
|
return v
|
|
|
|
|
|
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, external=None):
|
|
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))
|
|
nsets = ngraphs = ncerts = 0
|
|
lowered = {}
|
|
# The sets index.json lists, then the falsification sets of section 14.3. Each is
|
|
# re-lowered, its lowered graphs compared, and its certificate recomputed from
|
|
# them and compared. An index set's graphs each load (steps 1-8), so its verdict
|
|
# meets acc64; a falsification set's verdict fails the profile it claims.
|
|
listed = [(e, None) for e in json.load(open(os.path.join(root, "certificates", "index.json")))["sets"]]
|
|
listed += [(e, e["claim"]) for e in json.load(open(os.path.join(root, "certificates", "falsification.json")))["sets"]]
|
|
for e, claim in listed:
|
|
what = "falsification set" if claim else "set"
|
|
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")
|
|
try:
|
|
roles = lower_set(e, files)
|
|
except LowerRefuse:
|
|
raise ValueError("the lowering refused the set")
|
|
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") and not claim:
|
|
lowered[(FORM_FILE[e["lowering"]], e["name"], role)] = b
|
|
ngraphs += 1
|
|
nsets += 1
|
|
v = check_set_certificate(root, e, roles, write)
|
|
ncerts += 1
|
|
if claim is None and not v["acc64"]:
|
|
raise ValueError("a lowered graph does not load: the verdict fails acc64")
|
|
if claim is not None and v[claim]:
|
|
raise ValueError("its verdict meets %s, the profile it must fail" % claim)
|
|
except Refuse as ex:
|
|
failures += 1
|
|
print("FAIL %s %s (%s): a lowered graph is refused at load: %s" % (what, e.get("name"), e.get("capability"), ex.code))
|
|
except Exception as ex:
|
|
failures += 1
|
|
print("FAIL %s %s (%s): %s" % (what, e.get("name"), e.get("capability"), ex))
|
|
print("re-lowered %d sets, %d graphs; recomputed %d certificates byte for byte" % (nsets, ngraphs, ncerts))
|
|
failures += check_lower_vectors(root, lowered)
|
|
failures += check_external(root, external or {}, write)
|
|
return failures
|
|
|
|
|
|
def graph_meets(gj, p):
|
|
"""Section 11, `role_profiles`, read from a certificate's graph: its inputs' and
|
|
scalars' declared ranges meet profile p's tensor width, and every node fits p."""
|
|
tw = WIDTH[p]["tensor"]
|
|
ranges = [(dec(r["lo"]), dec(r["hi"])) for r in gj["inputs"] + gj["scalars"]]
|
|
return all(bits(r) <= tw for r in ranges) and all(p in nd["fits"] for nd in gj["nodes"])
|
|
|
|
|
|
def check_external(root, external, write):
|
|
"""The sets of external.json, whose certificates are committed although the sets
|
|
are not. With the set's files (`external` maps an entry's name to the directory
|
|
that holds them), the certificate is recomputed and compared. Without them, only
|
|
what needs no set: the certificate is canonical, names its set, its capability
|
|
and its lowering, its verdict meets the profile the entry gives, and each graph
|
|
the entry's `role_profiles` names meets its own profile (section 11)."""
|
|
failures = n_full = n_part = 0
|
|
for e in json.load(open(os.path.join(root, "certificates", "external.json")))["sets"]:
|
|
try:
|
|
want = open(cert_path(root, e), "rb").read()
|
|
d = e["name"] in external
|
|
if d:
|
|
files = {m: open(os.path.join(external[e["name"]], f), "rb").read() for m, f in e["members"].items()}
|
|
if set_digest(files) != e["set_sha256"]:
|
|
raise ValueError("the set's digest differs")
|
|
try:
|
|
roles = lower_set(e, files)
|
|
except LowerRefuse:
|
|
raise ValueError("the lowering refused the set")
|
|
check_set_certificate(root, e, roles, write)
|
|
c = json.loads(want)
|
|
if canonical(c) != want:
|
|
raise ValueError("the certificate is not canonical")
|
|
if (c["set_sha256"], c["capability"], c["lowering"]) != (e["set_sha256"], e["capability"], e["lowering"]):
|
|
raise ValueError("the certificate names another set")
|
|
if c["verdict"][e["profile"]] is not True:
|
|
raise ValueError("its verdict does not meet %s" % e["profile"])
|
|
for role, rp in sorted(e.get("role_profiles", {}).items()):
|
|
gj = [g for g in c["graphs"] if g["name"] == role]
|
|
if len(gj) != 1 or not graph_meets(gj[0], rp):
|
|
raise ValueError("its graph %s does not meet %s" % (role, rp))
|
|
if d:
|
|
n_full += 1
|
|
else:
|
|
n_part += 1
|
|
except Refuse as ex:
|
|
failures += 1
|
|
print("FAIL external set %s: a lowered graph is refused at load: %s" % (e["name"], ex.code))
|
|
except Exception as ex:
|
|
failures += 1
|
|
print("FAIL external set %s: %s" % (e["name"], ex))
|
|
print("external sets: %d certificates recomputed byte for byte, %d checked without their set" % (n_full, n_part))
|
|
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 certify_main(args):
|
|
"""`certify`: one set, or WHG1 graphs on their own, to a certificate (section 8.5),
|
|
with the exit status of the bounds tool: 0 the verdict meets the profile, and each
|
|
graph a --role-profile names meets its own (section 11, `role_profiles`); 1 one of
|
|
them does not (the certificate is still written); 2 a graph is malformed or
|
|
refused at load; 3 `--check` finds a difference; 4 the lowering refuses the set;
|
|
5 usage, a --role-profile role among them that names no graph."""
|
|
opts = {"--capability": None, "--lowering": None, "--c": None, "--profile": "acc64", "--out": None, "--check": None}
|
|
members, graphs, role_profiles = [], [], {}
|
|
k = 0
|
|
while k < len(args):
|
|
a = args[k]
|
|
if k + 1 >= len(args) or (a not in opts and a not in ("--member", "--whg1", "--role-profile")):
|
|
print(CERTIFY_USAGE, file=sys.stderr)
|
|
return 5
|
|
v = args[k + 1]
|
|
if a in ("--member", "--whg1", "--role-profile"):
|
|
if "=" not in v:
|
|
print(CERTIFY_USAGE, file=sys.stderr)
|
|
return 5
|
|
name, val = v.split("=", 1)
|
|
if a == "--role-profile":
|
|
if name in role_profiles or val not in PROFILES:
|
|
print(CERTIFY_USAGE, file=sys.stderr)
|
|
return 5
|
|
role_profiles[name] = val
|
|
else:
|
|
(members if a == "--member" else graphs).append((name, val))
|
|
else:
|
|
opts[a] = v
|
|
k += 2
|
|
if bool(members) == bool(graphs) or opts["--profile"] not in PROFILES or (members and not (opts["--capability"] and opts["--lowering"])):
|
|
print(CERTIFY_USAGE, file=sys.stderr)
|
|
return 5
|
|
try:
|
|
if graphs:
|
|
roles = [(name, open(path, "rb").read()) for name, path in graphs]
|
|
capability, lowering, digest = opts["--capability"], "whg1", None
|
|
else:
|
|
files = {name: open(path, "rb").read() for name, path in members}
|
|
capability, lowering, digest = opts["--capability"], opts["--lowering"], set_digest(files)
|
|
entry = {"capability": capability, "lowering": lowering, "c": opts["--c"], "set_sha256": digest}
|
|
try:
|
|
roles = lower_set(entry, files)
|
|
except (LowerRefuse, KeyError, ValueError) as ex:
|
|
print("the lowering refused the set%s" % (": %s" % ex if str(ex) else ""), file=sys.stderr)
|
|
return 4
|
|
except OSError as ex:
|
|
print(ex, file=sys.stderr)
|
|
return 5
|
|
unknown = sorted(set(role_profiles) - {name for name, _ in roles})
|
|
if unknown:
|
|
print("--role-profile names no graph of the set: %s" % ", ".join(unknown), file=sys.stderr)
|
|
print(CERTIFY_USAGE, file=sys.stderr)
|
|
return 5
|
|
try:
|
|
got, v, role_fails = certificate(capability, lowering, digest, roles, role_profiles)
|
|
except Refuse as ex:
|
|
print("a graph is refused at load: %s" % ex.code, file=sys.stderr)
|
|
return 2
|
|
if opts["--out"]:
|
|
open(opts["--out"], "wb").write(got)
|
|
else:
|
|
sys.stdout.write(got.decode("ascii"))
|
|
sys.stdout.flush()
|
|
p = opts["--profile"]
|
|
if opts["--check"]:
|
|
want = open(opts["--check"], "rb").read()
|
|
if got != want:
|
|
print("--check: %s" % first_difference(got, want), file=sys.stderr)
|
|
return 3
|
|
status = 0
|
|
if not v[p]:
|
|
print("the verdict does not meet %s: %s" % (p, cw(v[p + "_first_failure"])), file=sys.stderr)
|
|
status = 1
|
|
for name, rp in role_profiles.items():
|
|
if name in role_fails:
|
|
print("graph %s does not meet %s: %s" % (name, rp, cw(role_fails[name])), file=sys.stderr)
|
|
status = 1
|
|
else:
|
|
print("graph %s meets %s" % (name, rp), file=sys.stderr)
|
|
if status:
|
|
return status
|
|
print("the verdict meets %s: %s" % (p, cw(v)), file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
CERTIFY_USAGE = """usage:
|
|
wai_ops_check.py certify --capability CAP --lowering ID [--c N] --member NAME=PATH ...
|
|
[--profile acc32|acc64] [--role-profile ROLE=acc32|acc64 ...]
|
|
[--out FILE] [--check FILE]
|
|
wai_ops_check.py certify --whg1 ROLE=PATH ... [--capability CAP]
|
|
[--profile acc32|acc64] [--role-profile ROLE=acc32|acc64 ...]
|
|
[--out FILE] [--check FILE]"""
|
|
|
|
|
|
def main(argv):
|
|
if len(argv) >= 2 and argv[1] == "certify":
|
|
return certify_main(argv[2:])
|
|
if len(argv) < 2 or len(argv) % 2:
|
|
print(__doc__)
|
|
return 2
|
|
write, external = None, {}
|
|
for k in range(2, len(argv), 2):
|
|
a, v = argv[k], argv[k + 1]
|
|
if a == "--write":
|
|
write = v
|
|
elif a == "--external" and "=" in v:
|
|
name, d = v.split("=", 1)
|
|
external[name] = d
|
|
else:
|
|
print(__doc__)
|
|
return 2
|
|
failures = check(argv[1], write, external)
|
|
if failures:
|
|
print("%d failures" % failures)
|
|
return 1
|
|
print("ops-conformance: every vector file, every lowered graph and every certificate rewritten byte for byte")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv))
|