mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
Daniel Koepping
31 lines
1.8 KiB
Python
31 lines
1.8 KiB
Python
import pickle, sys
|
|
from tinygrad.helpers import getenv, is_image_shape, temp
|
|
from tinygrad.uop.ops import Ops
|
|
from tinygrad.viz.serve import VizData, _reconstruct, canonicalize_ast
|
|
|
|
with open(sys.argv[1] if len(sys.argv) > 1 else temp("rewrites.pkl", append_user=True), "rb") as f: data = VizData(pickle.load(f))
|
|
steps = [s for key, ctx in zip(data.trace.keys, data.trace.rewrites) if key.display_name.startswith("JIT ")
|
|
for s in ctx if s.name == "View captured linear"]
|
|
assert len(steps) == 1, f"expected one JIT compilation, found {len(steps)}"
|
|
asts = [call.without_after.src[0] for call in _reconstruct(data, steps[0].sink).src]
|
|
by_ast = {canonicalize_ast(_reconstruct(data, b.sink)): _reconstruct(data, p.sink) for ctx in data.trace.rewrites
|
|
for b in ctx if b.name == "View Base AST" for p in ctx if p.name == "View Program"}
|
|
programs = [by_ast[canonicalize_ast(ast)] if ast.op is Ops.SINK else ast for ast in asts if ast.op in {Ops.SINK, Ops.PROGRAM}]
|
|
read_image, gated_read_image = 0, 0
|
|
for prg in programs:
|
|
uops = prg.src[1].src
|
|
reads = {u for u in uops if u.op is Ops.LOAD and u.src[0].op is Ops.INDEX and is_image_shape(u.src[0].src[0]._shape)}
|
|
gated = {u for u in reads if len(u.src) == 3}
|
|
# include image channels masked after loading, count each load only once.
|
|
for u in uops:
|
|
if u.op is not Ops.WHERE: continue
|
|
for value in u.src[1:]:
|
|
while value.op in {Ops.INDEX, Ops.SHRINK, Ops.CAST, Ops.BITCAST}: value = value.src[0]
|
|
if value in reads: gated.add(value)
|
|
read_image += len(reads)
|
|
gated_read_image += len(gated)
|
|
|
|
for name, count in zip(("KERNEL_COUNT", "READ_IMAGE", "GATED_READ_IMAGE"), (len(programs), read_image, gated_read_image)):
|
|
print(f"{name}={count}")
|
|
if (allowed:=getenv(f"ALLOWED_{name}", -1)) != -1: assert count == allowed, f"ALLOWED_{name}={allowed}, got {count}"
|