mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-07 22:23:55 -05:00
replace mop arg with vectorized index (#12695)
* replace mop arg with vectorized index * tests passing * better viz * no compile4
This commit is contained in:
@@ -1,47 +0,0 @@
|
||||
import sys
|
||||
from tinygrad import Tensor, fetch, GlobalCounters, dtypes
|
||||
from tinygrad.uop.ops import UOp
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from tinygrad.schedule.rangeify import get_rangeify_map
|
||||
from tinygrad.engine.schedule import create_schedule_with_vars
|
||||
from tinygrad.engine.realize import run_schedule
|
||||
|
||||
# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
|
||||
|
||||
OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
|
||||
OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
|
||||
|
||||
if __name__ == "__main__":
|
||||
onnx_file = fetch(OPENPILOT_MODEL)
|
||||
run_onnx = OnnxRunner(onnx_file)
|
||||
|
||||
inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
|
||||
out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
|
||||
root = out.uop
|
||||
targets = [x.uop for x in inputs.values()]
|
||||
print(targets)
|
||||
|
||||
# TODO: abstract this from gradient?
|
||||
|
||||
# compute the target path (top down)
|
||||
in_target_path: dict[UOp, bool] = {}
|
||||
for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
|
||||
independent_set = {}
|
||||
for u in root.toposort():
|
||||
if in_target_path[u]:
|
||||
for s in u.src:
|
||||
if not in_target_path[s]:
|
||||
independent_set[s] = None
|
||||
independent = UOp.sink(*independent_set.keys())
|
||||
kernelized = get_rangeify_map(independent)
|
||||
independent = independent.substitute(kernelized)
|
||||
schedule, var_vals = create_schedule_with_vars(independent)
|
||||
run_schedule(schedule)
|
||||
|
||||
print("**** real ****")
|
||||
GlobalCounters.reset()
|
||||
out.uop = root.substitute(kernelized)
|
||||
out.kernelize()
|
||||
|
||||
# realize
|
||||
out.realize()
|
||||
Reference in New Issue
Block a user