Switch PC running to tinygrad with LLVM (#34660)

* try llvm on cpu

* update lock

* Make sure test fails

* With BEAM=3 its fast

* small fixes

* more fixes

* Fix macos

* do dm too

* rm ort

selfdrive/modeld/SConscript

@@ -41,8 +41,10 @@ lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_file
 pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
 if arch == 'larch64':
   device_string = 'QCOM=1'
-else:
+elif arch == 'Darwin':
   device_string = 'CLANG=1 IMAGE=0'
+else:
+  device_string = 'LLVM=1 LLVMOPT=1 BEAM=0 IMAGE=0'
 
 for model_name in ['supercombo', 'dmonitoring_model']:
   fn = File(f"models/{model_name}").abspath

selfdrive/modeld/dmonitoringmodeld.py

@@ -1,13 +1,13 @@
 #!/usr/bin/env python3
 import os
 from openpilot.system.hardware import TICI
+from tinygrad.tensor import Tensor
+from tinygrad.dtype import dtypes
 if TICI:
-  from tinygrad.tensor import Tensor
-  from tinygrad.dtype import dtypes
   from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address
   os.environ['QCOM'] = '1'
 else:
-  from openpilot.selfdrive.modeld.runners.ort_helpers import make_onnx_cpu_runner
+  os.environ['LLVM'] = '1'
 import math
 import time
 import pickle
@@ -34,7 +34,6 @@ OUTPUT_SIZE = 84 + FEATURE_LEN
 
 PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
-MODEL_PATH = Path(__file__).parent / 'models/dmonitoring_model.onnx'
 MODEL_PKL_PATH = Path(__file__).parent / 'models/dmonitoring_model_tinygrad.pkl'
 
 
@@ -78,12 +77,9 @@ class ModelState:
       'calib': np.zeros((1, CALIB_LEN), dtype=np.float32),
     }
 
-    if TICI:
-      self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
-      with open(MODEL_PKL_PATH, "rb") as f:
-        self.model_run = pickle.load(f)
-    else:
-      self.onnx_cpu_runner = make_onnx_cpu_runner(MODEL_PATH)
+    self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
+    with open(MODEL_PKL_PATH, "rb") as f:
+      self.model_run = pickle.load(f)
 
   def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]:
     self.numpy_inputs['calib'][0,:] = calib
@@ -96,12 +92,10 @@ class ModelState:
       if 'input_img' not in self.tensor_inputs:
         self.tensor_inputs['input_img'] = qcom_tensor_from_opencl_address(input_img_cl.mem_address, (1, MODEL_WIDTH*MODEL_HEIGHT), dtype=dtypes.uint8)
     else:
-      self.numpy_inputs['input_img'] = self.frame.buffer_from_cl(input_img_cl).reshape((1, MODEL_WIDTH*MODEL_HEIGHT))
+      self.tensor_inputs['input_img'] = Tensor(self.frame.buffer_from_cl(input_img_cl).reshape((1, MODEL_WIDTH*MODEL_HEIGHT)), dtype=dtypes.uint8).realize()
 
-    if TICI:
-      output = self.model_run(**self.tensor_inputs).numpy().flatten()
-    else:
-      output = self.onnx_cpu_runner.run(None, self.numpy_inputs)[0].flatten()
+
+    output = self.model_run(**self.tensor_inputs).numpy().flatten()
 
     t2 = time.perf_counter()
     return output, t2 - t1

selfdrive/modeld/modeld.py

@@ -1,15 +1,13 @@
 #!/usr/bin/env python3
 import os
 from openpilot.system.hardware import TICI
-
-#
+from tinygrad.tensor import Tensor
+from tinygrad.dtype import dtypes
 if TICI:
-  from tinygrad.tensor import Tensor
-  from tinygrad.dtype import dtypes
   from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address
   os.environ['QCOM'] = '1'
 else:
-  from openpilot.selfdrive.modeld.runners.ort_helpers import make_onnx_cpu_runner
+  os.environ['LLVM'] = '1'
 import time
 import pickle
 import numpy as np
@@ -37,7 +35,6 @@ from openpilot.selfdrive.modeld.models.commonmodel_pyx import DrivingModelFrame,
 PROCESS_NAME = "selfdrive.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 
-MODEL_PATH = Path(__file__).parent / 'models/supercombo.onnx'
 MODEL_PKL_PATH = Path(__file__).parent / 'models/supercombo_tinygrad.pkl'
 METADATA_PATH = Path(__file__).parent / 'models/supercombo_metadata.pkl'
 
@@ -78,12 +75,9 @@ class ModelState:
     self.output = np.zeros(net_output_size, dtype=np.float32)
     self.parser = Parser()
 
-    if TICI:
-      self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
-      with open(MODEL_PKL_PATH, "rb") as f:
-        self.model_run = pickle.load(f)
-    else:
-      self.onnx_cpu_runner = make_onnx_cpu_runner(MODEL_PATH)
+    self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
+    with open(MODEL_PKL_PATH, "rb") as f:
+      self.model_run = pickle.load(f)
 
   def slice_outputs(self, model_outputs: np.ndarray) -> dict[str, np.ndarray]:
     parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in self.output_slices.items()}
@@ -113,15 +107,14 @@ class ModelState:
           self.tensor_inputs[key] = qcom_tensor_from_opencl_address(imgs_cl[key].mem_address, self.input_shapes[key], dtype=dtypes.uint8)
     else:
       for key in imgs_cl:
-        self.numpy_inputs[key] = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key]).astype(dtype=np.float32)
+        self.numpy_inputs[key] = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key])
+        self.tensor_inputs[key] = Tensor(self.numpy_inputs[key], dtype=dtypes.uint8).realize()
+
 
     if prepare_only:
       return None
 
-    if TICI:
-      self.output = self.model_run(**self.tensor_inputs).numpy().flatten()
-    else:
-      self.output = self.onnx_cpu_runner.run(None, self.numpy_inputs)[0].flatten()
+    self.output = self.model_run(**self.tensor_inputs).numpy().flatten()
 
     outputs = self.parser.parse_outputs(self.slice_outputs(self.output))
 

selfdrive/modeld/runners/ort_helpers.py

@@ -1,36 +0,0 @@
-import onnx
-import onnxruntime as ort
-import numpy as np
-import itertools
-
-ORT_TYPES_TO_NP_TYPES = {'tensor(float16)': np.float16, 'tensor(float)': np.float32, 'tensor(uint8)': np.uint8}
-
-def attributeproto_fp16_to_fp32(attr):
-  float32_list = np.frombuffer(attr.raw_data, dtype=np.float16)
-  attr.data_type = 1
-  attr.raw_data = float32_list.astype(np.float32).tobytes()
-
-def convert_fp16_to_fp32(model):
-  for i in model.graph.initializer:
-    if i.data_type == 10:
-      attributeproto_fp16_to_fp32(i)
-  for i in itertools.chain(model.graph.input, model.graph.output):
-    if i.type.tensor_type.elem_type == 10:
-      i.type.tensor_type.elem_type = 1
-  for i in model.graph.node:
-    if i.op_type == 'Cast' and i.attribute[0].i == 10:
-      i.attribute[0].i = 1
-    for a in i.attribute:
-      if hasattr(a, 't'):
-        if a.t.data_type == 10:
-          attributeproto_fp16_to_fp32(a.t)
-  return model.SerializeToString()
-
-
-def make_onnx_cpu_runner(model_path):
-  options = ort.SessionOptions()
-  options.intra_op_num_threads = 4
-  options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
-  options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-  model_data = convert_fp16_to_fp32(onnx.load(model_path))
-  return ort.InferenceSession(model_data,  options, providers=['CPUExecutionProvider'])