From ad119af6a511373e1c016a6525ab733f14a60c51 Mon Sep 17 00:00:00 2001
From: Comma Device <device@comma.ai>
Date: Mon, 11 Nov 2024 15:58:27 -0500
Subject: [PATCH] compile fixes

---
 examples/openpilot/compile3.py | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/examples/openpilot/compile3.py b/examples/openpilot/compile3.py
index 69331a6d..c781d7cb 100644
--- a/examples/openpilot/compile3.py
+++ b/examples/openpilot/compile3.py
@@ -5,8 +5,8 @@ if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
 if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
-from tinygrad import fetch, Tensor, TinyJit, Device, Context, GlobalCounters
-from tinygrad.helpers import OSX, DEBUG, getenv
+from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters
+from tinygrad.helpers import DEBUG, getenv
 from tinygrad.tensor import _from_np_dtype
 
 import onnx
@@ -14,14 +14,14 @@ from onnx.helper import tensor_dtype_to_np_dtype
 from extra.onnx import get_run_onnx   # TODO: port to main tinygrad
 
 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 = "/tmp/openpilot.pkl"
+OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
 
 def compile():
   # hack to fix GPU on OSX: max doesn't work on half, see test/external/external_gpu_fail_osx.py
-  if OSX:
-    from tinygrad.ops import BinaryOps
-    from tinygrad.renderer.cstyle import ClangRenderer, CStyleLanguage
-    CStyleLanguage.code_for_op[BinaryOps.MAX] = ClangRenderer.code_for_op[BinaryOps.MAX]
+  #if OSX:
+  #  from tinygrad.ops import BinaryOps
+  #  from tinygrad.renderer.cstyle import ClangRenderer, CStyleLanguage
+  #  CStyleLanguage.code_for_op[BinaryOps.MAX] = ClangRenderer.code_for_op[BinaryOps.MAX]
 
   Tensor.no_grad = True
   Tensor.training = False
@@ -32,7 +32,12 @@ def compile():
   print("loaded model")
 
   input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
-  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
+  input_types = {inp.name: np.float32 for inp in onnx_model.graph.input}
+  if 'input_img' in input_shapes:
+    input_types['input_img'] = np.uint8
+  else:
+    input_types['input_imgs'] = np.uint8
+    input_types['big_input_imgs'] = np.uint8
   Tensor.manual_seed(100)
   new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in sorted(input_shapes.items())}
   print("created tensors")
@@ -65,7 +70,11 @@ def test(test_val=None):
   new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
                 sorted(zip(run.captured.expected_names, run.captured.expected_st_vars_dtype_device))}
   for _ in range(20):
+    inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
     st = time.perf_counter()
+    for k in new_inputs:
+      if 'img' not in k: # dont need to init img tensors, those are backed by openCL GPU memory
+        new_inputs[k] = Tensor(inputs_numpy[k])
     out = run(**new_inputs)
     mt = time.perf_counter()
     val = out['outputs'].numpy()