Vehicle Researcher
6adb63b915
date: 2026-06-04T09:49:56 master commit: c0ab3550eca2e9daf197c46b7e4b24aa9637cf2e
35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
# benchmark speed of pyrender for all created UOps saved with TRACK_MATCH_STATS=2
|
|
import functools, pickle
|
|
from tinygrad.uop.ops import UOp, Ops
|
|
from tinygrad.helpers import tqdm, temp, time_to_str, cpu_profile
|
|
|
|
BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}
|
|
|
|
@functools.cache
|
|
def create_uop(a:int) -> UOp:
|
|
op, dtype, src, arg, *rest = trace.uop_fields[a]
|
|
return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)
|
|
|
|
if __name__ == "__main__":
|
|
# load rewrite trace
|
|
with open(temp("rewrites.pkl", append_user=True), "rb") as f:
|
|
trace = pickle.load(f)
|
|
|
|
# benchmark
|
|
result:list[tuple[str, int]] = []
|
|
try:
|
|
for steps in tqdm(trace.rewrites):
|
|
for r in steps:
|
|
for _,yn,_,__ in r.matches:
|
|
y = create_uop(yn)
|
|
if y.op in BENCHMARK_OPS:
|
|
with cpu_profile("pyrender") as e:
|
|
try: ren = y.render()
|
|
except Exception: ren = "PYRENDER_ERR"
|
|
result.append((ren, float(e.en-e.st)/1e6))
|
|
finally:
|
|
N = 10
|
|
print(f"Slowst {N} renders from {len(result)} samples:")
|
|
for ren,tm in sorted(result, key=lambda x:x[1], reverse=True)[:N]:
|
|
print(f"{time_to_str(tm).strip():<10s} {ren}")
|