From 577a0b4cfa38ed4c9b8055ef6c81fa8a46c59716 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 22 May 2025 10:47:34 -0700 Subject: [PATCH] openpilot compile4 (wip) (#10407) * openpilot compile4 * add copies * remove junk --- examples/openpilot/compile4.py | 49 ++++++++++++++++++++++++++++++++++ extra/onnx.py | 3 +++ tinygrad/runtime/ops_npy.py | 1 + tinygrad/tensor.py | 6 ++--- tinygrad/uop/__init__.py | 1 + tinygrad/uop/ops.py | 2 +- 6 files changed, 58 insertions(+), 4 deletions(-) create mode 100644 examples/openpilot/compile4.py diff --git a/examples/openpilot/compile4.py b/examples/openpilot/compile4.py new file mode 100644 index 0000000000..695b620322 --- /dev/null +++ b/examples/openpilot/compile4.py @@ -0,0 +1,49 @@ +import sys, onnx +from tinygrad import Tensor, fetch, GlobalCounters +from tinygrad.uop import UOp +from tinygrad.frontend.onnx import OnnxRunner +from tinygrad.engine.grouper import get_becomes_map +from tinygrad.engine.schedule import create_schedule_with_vars +from tinygrad.engine.realize import run_schedule + +# NOLOCALS=1 GPU=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__": + fn = fetch(OPENPILOT_MODEL) + onnx_file = fetch(OPENPILOT_MODEL) + onnx_model = onnx.load(onnx_file) + run_onnx = OnnxRunner(onnx_model) + + inputs = run_onnx.get_empty_input_data("npy") + out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu') + root = out.lazydata + targets = [x.lazydata 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_becomes_map(independent) + independent = independent.substitute(kernelized) + schedule, var_vals, becomes_map = create_schedule_with_vars(independent) + run_schedule(schedule) + + print("**** real ****") + GlobalCounters.reset() + out.lazydata = root.substitute(kernelized).substitute(becomes_map) + out.kernelize() + + # realize + out.realize() diff --git a/extra/onnx.py b/extra/onnx.py index 11a8623fc2..6cdd4d6e03 100644 --- a/extra/onnx.py +++ b/extra/onnx.py @@ -151,6 +151,9 @@ class OnnxRunner: return real_fxn(*inps, **opts) raise NotImplementedError(f"{op=} not supported") + def get_empty_input_data(self, device:str|None=None) -> dict[str, Tensor]: + return {name:Tensor.empty(*spec.shape, device=device, dtype=spec.dtype) for name, spec in self.graph_inputs.items()} + def __call__(self, inputs:dict[str, Any], debug=debug): for name, input_spec in self.graph_inputs.items(): if name not in inputs: raise RuntimeError(f"Please provide input data for {name}") diff --git a/tinygrad/runtime/ops_npy.py b/tinygrad/runtime/ops_npy.py index b0e2edf750..f40bfcdeeb 100644 --- a/tinygrad/runtime/ops_npy.py +++ b/tinygrad/runtime/ops_npy.py @@ -3,6 +3,7 @@ from tinygrad.helpers import flat_mv from tinygrad.device import Compiled, Allocator class NpyAllocator(Allocator['NpyDevice']): + def _alloc(self, size:int, options=None) -> np.ndarray: return np.empty(size, dtype=np.uint8) def _as_buffer(self, src:np.ndarray) -> memoryview: return flat_mv(np.require(src, requirements='C').data) def _copyout(self, dest:memoryview, src:np.ndarray): dest[:] = self._as_buffer(src) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 6e54caca72..2603906fd9 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -19,6 +19,8 @@ from tinygrad.engine.grouper import get_becomes_map # *** all in scope Tensors are here. this gets relevant UOps *** all_tensors: set[weakref.ref[Tensor]] = set() +def _find_all_tensors_for_uops(all_uops: set[UOp]) -> list[Tensor]: + return [t for tref in all_tensors if (t:=tref()) is not None and t.lazydata in all_uops] def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str|None=None) -> None: # get all children of keys in applied_map @@ -32,9 +34,7 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str|None=None) -> Non # link the found UOps back to Tensors. exit early if there's no Tensors to realize # NOTE: this uses all_tensors, but it's fast - fixed_tensors: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and t.lazydata in all_uops] - - if len(fixed_tensors): + if len(fixed_tensors := _find_all_tensors_for_uops(all_uops)): # potentially rewrite all the discovered Tensors sink = UOp.sink(*[t.lazydata for t in fixed_tensors]) new_sink = sink.substitute(applied_map, name=name) diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index e69de29bb2..efae111520 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -0,0 +1 @@ +from tinygrad.uop.ops import UOp, Ops # noqa: F401 \ No newline at end of file diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index e7a2bca240..86ee419bc0 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -582,7 +582,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return bound_vars.union(set([x for x in all_vars if x not in bound_var_base])) def variables(self) -> list[Variable]: st_vars: list[set[Variable]] = [x.st_arg.vars() for x in self.toposort() if x.op in GroupOp.Buffer] - return sorted(set.union(*st_vars, [x.unbind()[0] if x.op is not Ops.DEFINE_VAR else x for x in self.vars()]), key=lambda v: v.arg) + return sorted(set.union(*st_vars, set([x.unbind()[0] if x.op is not Ops.DEFINE_VAR else x for x in self.vars()])), key=lambda v: v.arg) # *** uop symbolic stuff ***