openpilot compile4 (wip) (#10407)

* openpilot compile4

* add copies

* remove junk

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()

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}")

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)
 

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)

tinygrad/uop/__init__.py

@@ -0,0 +1 @@
+from tinygrad.uop.ops import UOp, Ops  # noqa: F401

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 ***