d1-2B-LiteRT / conversion /make_rows.py
mlboydaisuke's picture
d1-2B for LiteRT: decision graphs at seven lengths, the eight-sequence and eight-signature files, seven picture towers, five SM8850 NPU files, the PyTorch-free host, the fixtures and the conversion scripts
df44977 verified
Raw History Blame Contribute Delete
15.5 kB
"""The device rows for d1-2B (the phones' test app reads them in its generic mode; the app is not published), made by
the PyTorch-free host (host/d1_host.py) with recording backends (replay.py): every graph call the host makes for a
request, laid out as the LiteRT backends lay it out (host/lrt_backend.py), becomes a row of the bucket's / canvas's
rows file.
$EXPORT conversion/make_rows.py static --set conversion/set_ids.json --out-dir out/rows
tower rows per canvas (twG<GH>x<GW>.json), the text requests' decide rows per bucket (decL<L>.json), the timing
rows of the card's workloads (tmL<L>.json, tmG<GH>x<GW>.json) and pic_index.json (the picture requests: the
canvases their pictures and the buckets their sequences take), SHA256SUMS
$EXPORT conversion/make_rows.py pic --rows-dir <rows> --round-dir <R> --src-tag <a tower leg> --bucket L \
--name <rows file name> --out-dir <dir> [--sha-file <file>]
the picture requests' decide rows of bucket L, their embeddings holding the tower output of the round's tower
legs (every DONE leg in R whose tag starts like --src-tag up to "tw": one per canvas), so the phone's decision
graph reads the phone's own picture features
A picture request's decide rows are never static: its embeddings hold the tower's output. The timing rows are the one
exception: they time the graph's work, which does not depend on the values, so their picture positions hold zeros
(written in the file as `picture_features: zeros`).
Run with $EXPORT (numpy, tokenizers, Pillow; no torch).
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import sys
import time
HERE = os.path.dirname(os.path.abspath(__file__))
WORK = os.path.normpath(os.path.join(HERE, ".."))
sys.path.insert(0, os.path.join(WORK, "host")) # the host modules (replay.py is in this folder)
sys.path.insert(0, os.path.join(WORK, "host"))
SNAP = os.path.join(WORK, "snapshot")
EMBED = os.path.join(WORK, "tables/embed_table.npy")
RECORD_FILES = ["set/requests.json", "set/red_arms.json", "set/extra_g96x72.json"]
MEDIA = "set/media_all.json"
WORKLOADS = {"one_question": ["wl_one_question"], "3_questions": ["wl_3_questions"], "5k_state": ["wl_5k_state"],
"384px": ["wl_384px"], "box": ["wl_box"], "boxes": ["wl_boxes"],
"64_states": [f"wl_64_states_{i:02d}" for i in range(64)]} # conversion/bench_host.py
def load_records():
out = {}
for p in RECORD_FILES:
doc = json.load(open(os.path.join(WORK, p)))
for r in (doc if isinstance(doc, list) else doc["records"]):
out[r["id"]] = r
return out
def open_media(media, ids):
from PIL import Image
items = []
for mid in ids or []:
m = media[mid]
p = os.path.join(WORK, m["path"])
data = open(p, "rb").read()
if m.get("sha256") and hashlib.sha256(data).hexdigest() != m["sha256"]:
raise SystemExit(f"media {mid}: sha256 mismatch")
items.append(Image.open(p))
return items
def run_records(host, dec, tow, recs, ids, media):
"""Each request through the host once (run_gate.py's path); the recording backends keep the calls under the id."""
for rid in ids:
req = recs[rid]["request"]
dec.tag = tow.tag = rid
n0 = len(dec.calls)
items = open_media(media, req.get("media"))
named = host.named(req["questions"])
host.run([(req.get("state"), list(named.values()), items)])
for n, c in enumerate(dec.calls[n0:]):
c["call"] = n
def grouped(calls, key):
out = {}
for c in calls:
out.setdefault(key(c), []).append(c)
return out
def sums_append(path, pairs):
if not path:
return
with open(path, "a") as fh:
for s, f in pairs:
fh.write(f"{s} {f}\n")
def cmd_static(a):
from d1_host import D1Host
from replay import (DECIDE_SIG, RecordingDecide, RecordingTower, decide_spec, tower_spec, write_rows)
recs = load_records()
media = json.load(open(os.path.join(WORK, MEDIA)))
sel = json.load(open(os.path.join(WORK, a.set)))
ids = sel["records"] + sel["extra_canvas"]
dec, tow = RecordingDecide(), RecordingTower()
host = D1Host(SNAP, dec, tow, embed_table=EMBED)
t0 = time.time()
run_records(host, dec, tow, recs, ids, media)
pic_ids = sorted({c["tag"] for c in tow.calls})
os.makedirs(a.out_dir, exist_ok=True)
written = []
# tower rows per canvas (every picture of the set)
for c, calls in sorted(grouped(tow.calls, lambda c: c["canvas"]).items()):
stem = f"twG{c[0]}x{c[1]}"
ins, outs = tower_spec(c, stem)
rows = [{"key": f"{x['tag']}/p{j}", "record": x["tag"], "grid": list(x["grid"])}
for j, x in enumerate(calls)]
written += write_rows(os.path.join(a.out_dir, stem + ".json"), "d1-2b_tower", "tower", ins, outs, rows,
[x["arrays"] for x in calls], {"made_by": "conversion/make_rows.py static"})
# decide rows of the text requests per bucket
text = [c for c in dec.calls if c["tag"] not in pic_ids]
for L, calls in sorted(grouped(text, lambda c: c["L"]).items()):
stem = f"decL{L}"
ins, outs = decide_spec(L, stem)
rows = [{"key": f"{x['tag']}/c{x['call']}", "record": x["tag"], "call": x["call"], "T": x["T"],
"state_len": x["state_len"]} for x in calls]
written += write_rows(os.path.join(a.out_dir, stem + ".json"), "d1-2b_decide", DECIDE_SIG[L], ins, outs,
rows, [x["arrays"] for x in calls], {"made_by": "conversion/make_rows.py static"})
pic = {}
for x in dec.calls:
if x["tag"] in pic_ids:
pic.setdefault(x["tag"], {"buckets": [], "canvases": []})["buckets"].append(x["L"])
for x in tow.calls:
pic[x["tag"]]["canvases"].append(list(x["canvas"]))
json.dump({"records": pic, "made_by": "conversion/make_rows.py static"},
open(os.path.join(a.out_dir, "pic_index.json"), "w"), indent=1)
# the card's workloads (timing; picture positions hold zeros: the tower is timed on its own canvas)
wdec, wtow = RecordingDecide(), RecordingTower()
whost = D1Host(SNAP, wdec, wtow, embed_table=EMBED)
wl_ids = [r for v in WORKLOADS.values() for r in v]
run_records(whost, wdec, wtow, recs, wl_ids, media)
by_rec = {wl: [r for r in v] for wl, v in WORKLOADS.items()}
sets_dec, sets_tow = {}, {}
for wl, rids in by_rec.items():
dcalls = [c for c in wdec.calls if c["tag"] in rids]
tcalls = [c for c in wtow.calls if c["tag"] in rids]
for c in dcalls:
sets_dec.setdefault(c["L"], {}).setdefault(wl, []).append(c)
for c in tcalls:
sets_tow.setdefault(c["canvas"], {}).setdefault(wl, []).append(c)
for L, by_wl in sorted(sets_dec.items()):
stem = f"tmL{L}"
ins, outs = decide_spec(L, stem)
calls = [c for wl in by_wl for c in by_wl[wl]]
rows = [{"key": f"{x['tag']}/c{x['call']}", "record": x["tag"], "call": x["call"], "T": x["T"]} for x in calls]
sets = [{"name": wl, "kind": "request" if len(v) > 1 else "single",
"rows": [f"{x['tag']}/c{x['call']}" for x in v]} for wl, v in by_wl.items()]
written += write_rows(os.path.join(a.out_dir, stem + ".json"), "d1-2b_decide", DECIDE_SIG[L], ins, outs, rows,
[x["arrays"] for x in calls], {"sets": sets, "picture_features": "zeros",
"made_by": "conversion/make_rows.py static"})
for c, by_wl in sorted(sets_tow.items()):
stem = f"tmG{c[0]}x{c[1]}"
ins, outs = tower_spec(c, stem)
calls = [x for wl in by_wl for x in by_wl[wl]]
rows = [{"key": f"{x['tag']}/p{j}", "record": x["tag"], "grid": list(x["grid"])} for j, x in enumerate(calls)]
keys = {id(x): r["key"] for x, r in zip(calls, rows)}
sets = [{"name": wl, "kind": "single", "rows": [keys[id(x)] for x in v]} for wl, v in by_wl.items()]
written += write_rows(os.path.join(a.out_dir, stem + ".json"), "d1-2b_tower", "tower", ins, outs, rows,
[x["arrays"] for x in calls], {"sets": sets, "made_by": "conversion/make_rows.py static"})
with open(os.path.join(a.out_dir, "SHA256SUMS"), "w") as fh:
for s, f in sorted(written, key=lambda t: t[1]):
fh.write(f"{s} {f}\n")
summary = {"records": len(ids), "picture_records": pic_ids, "tower_calls": len(tow.calls),
"decide_calls": len(dec.calls), "text_decide_calls": len(text),
"per_bucket_text": {str(L): len(v) for L, v in sorted(grouped(text, lambda c: c["L"]).items())},
"per_bucket_picture": {str(L): len(v) for L, v in sorted(grouped(
[c for c in dec.calls if c["tag"] in pic_ids], lambda c: c["L"]).items())},
"per_canvas": {f"{c[0]}x{c[1]}": len(v) for c, v in sorted(grouped(tow.calls, lambda c: c["canvas"]).items())},
"timing_sets": {f"L{L}": {wl: len(v) for wl, v in by_wl.items()} for L, by_wl in sets_dec.items()},
"timing_tower_sets": {f"G{c[0]}x{c[1]}": {wl: len(v) for wl, v in by_wl.items()}
for c, by_wl in sets_tow.items()},
"files": len(written), "bytes": sum(os.path.getsize(os.path.join(a.out_dir, f)) for _, f in written),
"wall_s": round(time.time() - t0, 1)}
json.dump(summary, open(os.path.join(a.out_dir, "static_summary.json"), "w"), indent=1)
print(json.dumps(summary))
def tower_sources(rows_dir, round_dir, prefix):
"""The round's DONE tower legs whose tag starts with `prefix` (one per canvas): GenericRows with their out file."""
from replay import GenericRows
out = {}
for name in sorted(os.listdir(round_dir)):
if not (name.startswith(prefix) and name.endswith(".json")) or name.endswith(".verdict"):
continue
tag = name[:-5]
rep = os.path.join(round_dir, name)
try:
doc = json.load(open(rep))
except Exception:
continue
if doc.get("status") != "DONE" or not str(doc.get("rows_file", "")).startswith("twG"):
continue
o = os.path.join(round_dir, f"out_{tag}.f32")
if not os.path.exists(o):
continue
rows = doc["rows_file"]
limit = (doc.get("summary") or {}).get("count") or 0
src = GenericRows(os.path.join(rows_dir, rows), o, rows_limit=limit)
if limit != len(src.rows):
raise SystemExit(f"tower leg {tag}: {limit} of {len(src.rows)} rows ran (stopped early?)")
out[rows] = (tag, src)
return out
def cmd_pic(a):
from d1_host import D1Host
from replay import DECIDE_SIG, RecordingDecide, RecordingTower, ReplayTower, decide_spec, write_rows
t0 = time.time()
pi = json.load(open(os.path.join(a.rows_dir, "pic_index.json")))["records"]
want = sorted(r for r, v in pi.items() if a.bucket in v["buckets"])
if not want:
raise SystemExit(f"no picture request has a sequence in bucket {a.bucket}")
prefix = a.src_tag[:a.src_tag.index("tw") + 2] if "tw" in a.src_tag else a.src_tag
srcs = tower_sources(a.rows_dir, a.round_dir, prefix)
need = sorted({f"twG{c[0]}x{c[1]}.json" for r in want for c in pi[r]["canvases"]})
missing = [n for n in need if n not in srcs]
dropped = []
if missing and not a.allow_missing:
raise SystemExit(f"bucket {a.bucket}: no DONE tower leg of prefix {prefix} for {missing} in {a.round_dir}")
if missing: # the requests whose pictures no tower leg ran leave the file (the replay then refuses them, by id)
dropped = [r for r in want if any(f"twG{c[0]}x{c[1]}.json" in missing for c in pi[r]["canvases"])]
want = [r for r in want if r not in dropped]
need = [n for n in need if n not in missing]
print(f"bucket {a.bucket}: no tower leg for {missing}: dropped {dropped}", file=sys.stderr)
replay = ReplayTower([srcs[n][1] for n in need]) if need else None
recs = load_records()
media = json.load(open(os.path.join(WORK, MEDIA)))
dec, tow = RecordingDecide(), RecordingTower(replay)
host = D1Host(SNAP, dec, tow, embed_table=EMBED)
run_records(host, dec, tow, recs, want, media)
calls = [c for c in dec.calls if c["L"] == a.bucket]
stem = a.name[:-5] if a.name.endswith(".json") else a.name
ins, outs = decide_spec(a.bucket, stem)
rows = [{"key": f"{x['tag']}/c{x['call']}", "record": x["tag"], "call": x["call"], "T": x["T"],
"state_len": x["state_len"]} for x in calls]
arrays = [x["arrays"] for x in calls]
n_text = 0
if a.with_text: # the bucket's text rows first (the static decL<L>.json, its slices read and checked)
import numpy as np
from replay import GenericRows
tp = os.path.join(a.rows_dir, f"decL{a.bucket}.json")
if os.path.exists(tp):
t = GenericRows(tp)
t_rows, t_arr = [], []
for r in t.rows:
d = {}
for spec in t.inputs:
n = int(np.prod(spec["shape"]))
mm = np.memmap(os.path.join(t.dir, spec["file"]), dtype="<f4", mode="r")
d[spec["name"]] = np.asarray(mm[r["index"] * n:(r["index"] + 1) * n]).reshape(spec["shape"])
t_arr.append(d)
t_rows.append({k: v for k, v in r.items() if k not in ("index", "sha256")})
rows, arrays, n_text = t_rows + rows, t_arr + arrays, len(t_rows)
os.makedirs(a.out_dir, exist_ok=True)
written = write_rows(os.path.join(a.out_dir, stem + ".json"), "d1-2b_decide", DECIDE_SIG[a.bucket], ins, outs,
rows, arrays,
{"picture_features": {n: srcs[n][0] for n in need}, "text_rows_first": n_text,
"pictures_missing": missing, "requests_dropped": dropped,
"round_dir": a.round_dir, "made_by": "conversion/make_rows.py pic"})
sums_append(a.sha_file, written)
print(json.dumps({"bucket": a.bucket, "name": stem + ".json", "rows": len(rows), "text_rows": n_text, "records": want,
"tower_legs": {n: srcs[n][0] for n in need}, "replayed": len(replay.used) if replay else 0, "dropped": dropped,
"wall_s": round(time.time() - t0, 1)}))
def main():
ap = argparse.ArgumentParser()
sub = ap.add_subparsers(dest="cmd", required=True)
s = sub.add_parser("static")
s.add_argument("--set", default="conversion/set_ids.json")
s.add_argument("--out-dir", default=os.path.join(WORK, "out/rows"))
p = sub.add_parser("pic")
p.add_argument("--rows-dir", required=True)
p.add_argument("--round-dir", required=True)
p.add_argument("--src-tag", required=True)
p.add_argument("--bucket", type=int, required=True)
p.add_argument("--name", required=True)
p.add_argument("--out-dir", required=True)
p.add_argument("--sha-file", default="")
p.add_argument("--with-text", action="store_true", help="the bucket's static text rows first, in the same file")
p.add_argument("--allow-missing", action="store_true", help="drop the requests whose canvas no tower leg ran")
a = ap.parse_args()
{"static": cmd_static, "pic": cmd_pic}[a.cmd](a)
if __name__ == "__main__":
main()