github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: just use compile3 (#12849)

Browse Source
This commit is contained in:
wozeparrot
2025-10-21 07:56:50 -07:00
committed by GitHub
parent c7336c3e31
commit 62e7b8b870
2 changed files with 6 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
test/external/external_benchmark_openpilot.py vendored
View File

@@ -1,63 +0,0 @@
import time, sys, hashlib
from pathlib import Path
from tinygrad.nn.onnx import OnnxRunner
from tinygrad import Tensor, dtypes, TinyJit
from tinygrad.helpers import IMAGE, GlobalCounters, fetch, colored, getenv, trange
import numpy as np
from extra.bench_log import BenchEvent, WallTimeEvent
OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"
if __name__ == "__main__":
run_onnx = OnnxRunner(fetch(OPENPILOT_MODEL))
Tensor.manual_seed(100)
input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
new_inputs_junk = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
new_inputs_junk_numpy = {k:v.numpy() for k,v in new_inputs_junk.items()}
# benchmark
for _ in range(5):
GlobalCounters.reset()
st = time.perf_counter_ns()
ret = next(iter(run_onnx(new_inputs_junk).values())).cast(dtypes.float32).numpy()
print(f"unjitted: {(time.perf_counter_ns() - st)*1e-6:7.4f} ms")
# NOTE: the inputs to a JIT must be first level arguments
run_onnx_jit = TinyJit(lambda **kwargs: run_onnx(kwargs), prune=True)
step_times = []
for _ in range(20):
GlobalCounters.reset()
st = time.perf_counter_ns()
with WallTimeEvent(BenchEvent.STEP):
# Need to cast non-image inputs from numpy, this is only realistic way to run model
inputs = {**{k:v for k,v in new_inputs_junk.items() if 'img' in k},
**{k:Tensor(v) for k,v in new_inputs_junk_numpy.items() if 'img' not in k}}
ret = next(iter(run_onnx_jit(**inputs).values())).cast(dtypes.float32).numpy()
step_times.append(t:=(time.perf_counter_ns() - st)*1e-6)
print(f"jitted: {t:7.4f} ms")
suffix = ""
if IMAGE.value < 2: suffix += f"_image{IMAGE.value}" # image=2 has no suffix for compatibility
if getenv("FLOAT16") == 1: suffix += "_float16"
path = Path(__file__).parent / "openpilot" / f"{hashlib.md5(OPENPILOT_MODEL.encode()).hexdigest()}{suffix}.npy"
# validate if we have records
tinygrad_out = next(iter(run_onnx_jit(**new_inputs).values())).cast(dtypes.float32).numpy()
if getenv("SAVE_OUTPUT"):
np.save(path, tinygrad_out)
print(f"saved output to {path}!")
elif getenv("FUZZ") and path.exists():
known_good_out = np.load(path)
for _ in trange(1000):
ret = next(iter(run_onnx_jit(**new_inputs).values())).cast(dtypes.float32).numpy()
np.testing.assert_allclose(known_good_out, ret, atol=1e-2, rtol=1e-2)
print(colored("fuzz validated!", "green"))
elif path.exists():
known_good_out = np.load(path)
np.testing.assert_allclose(known_good_out, tinygrad_out, atol=1e-2, rtol=1e-2)
print(colored("outputs validated!", "green"))
else:
print(colored("skipping validation", "yellow"))