Squash

repro.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+while :
+do
+  rm selfdrive/modeld/models/supercombo_tinygrad.pkl
+  ./system/manager/build.py 
+done
+

selfdrive/modeld/SConscript

@@ -13,15 +13,6 @@ common_src = [
   "transforms/transform.cc",
 ]
 
-thneed_src_common = [
-  "thneed/thneed_common.cc",
-  "thneed/serialize.cc",
-]
-
-thneed_src_qcom = thneed_src_common + ["thneed/thneed_qcom2.cc"]
-thneed_src_pc = thneed_src_common + ["thneed/thneed_pc.cc"]
-thneed_src = thneed_src_qcom if arch == "larch64" else thneed_src_pc
-
 # SNPE except on Mac and ARM Linux
 snpe_lib = []
 if arch != "Darwin" and arch != "aarch64":
@@ -59,16 +50,13 @@ fn = File("models/supercombo").abspath
 cmd = f'python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
 lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files, cmd)
 
-# Build thneed model
-if arch == "larch64" or GetOption('pc_thneed'):
-  tinygrad_opts = []
-  if not GetOption('pc_thneed'):
-    # use FLOAT16 on device for speed + don't cache the CL kernels for space
-    tinygrad_opts += ["FLOAT16=1", "PYOPENCL_NO_CACHE=1"]
-  cmd = f"cd {Dir('#').abspath}/tinygrad_repo && " + ' '.join(tinygrad_opts) + f" python3 openpilot/compile2.py {fn}.onnx {fn}.thneed"
-
-  lenv.Command(fn + ".thneed", [fn + ".onnx"] + tinygrad_files, cmd)
-
-  thneed_lib = env.SharedLibrary('thneed', thneed_src, LIBS=[gpucommon, common, 'OpenCL', 'dl'])
-  thneedmodel_lib = env.Library('thneedmodel', ['runners/thneedmodel.cc'])
-  lenvCython.Program('runners/thneedmodel_pyx.so', 'runners/thneedmodel_pyx.pyx', LIBS=envCython["LIBS"]+[thneedmodel_lib, thneed_lib, gpucommon, common, 'dl', 'OpenCL'])
+# Compile tinygrad model
+# TODO this is all super hacky
+pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
+if arch == 'larch64':
+  device_string = 'QCOM=1'
+else:
+  device_string = 'GPU=1'
+fn = File("models/supercombo").abspath
+cmd = f'{pythonpath_string} {device_string} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx && mv /tmp/openpilot.pkl {fn}_tinygrad.pkl'
+lenv.Command(fn + "_tinygrad.pkl", [fn + ".onnx"] + tinygrad_files, cmd)

selfdrive/modeld/dmonitoringmodeld

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+
+DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null && pwd)"
+exec "$DIR/dmonitoringmodeld.py" "$@"

selfdrive/modeld/modeld.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 import os
+## TODO this is hack
+os.environ['QCOM'] = '1'
 import time
 import pickle
 import numpy as np
@@ -19,19 +21,21 @@
 from openpilot.system import sentry
 from openpilot.selfdrive.car.card import convert_to_capnp
 from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
-from openpilot.selfdrive.modeld.runners import ModelRunner, Runtime
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
 from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState
 from openpilot.selfdrive.modeld.constants import ModelConstants
 from openpilot.selfdrive.modeld.models.commonmodel_pyx import ModelFrame, CLContext
 
+from tinygrad.tensor import Tensor
+Tensor.manual_seed(1337)
+Tensor.no_grad = True
+
 PROCESS_NAME = "selfdrive.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 
-MODEL_PATHS = {
-  ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed',
-  ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'}
 
+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'
 
 MODEL_WIDTH = 512
@@ -53,20 +57,19 @@ class ModelState:
   inputs: dict[str, np.ndarray]
   output: np.ndarray
   prev_desire: np.ndarray  # for tracking the rising edge of the pulse
-  model: ModelRunner
 
   def __init__(self, context: CLContext):
     self.frame = ModelFrame(context)
     self.wide_frame = ModelFrame(context)
     self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)
     self.inputs = {
-      'desire': np.zeros(ModelConstants.DESIRE_LEN * (ModelConstants.HISTORY_BUFFER_LEN+1), dtype=np.float32),
-      'traffic_convention': np.zeros(ModelConstants.TRAFFIC_CONVENTION_LEN, dtype=np.float32),
-      'lateral_control_params': np.zeros(ModelConstants.LATERAL_CONTROL_PARAMS_LEN, dtype=np.float32),
-      'prev_desired_curv': np.zeros(ModelConstants.PREV_DESIRED_CURV_LEN * (ModelConstants.HISTORY_BUFFER_LEN+1), dtype=np.float32),
-      'features_buffer': np.zeros(ModelConstants.HISTORY_BUFFER_LEN * ModelConstants.FEATURE_LEN, dtype=np.float32),
-      'input_imgs': None,
-      'big_input_imgs': None,
+      'input_imgs': np.zeros((1, 12, 128, 256), dtype=np.float16),
+      'big_input_imgs': np.zeros((1, 12, 128,  256), dtype=np.float16),
+      'desire': np.zeros((1, (ModelConstants.HISTORY_BUFFER_LEN+1), ModelConstants.DESIRE_LEN), dtype=np.float16),
+      'traffic_convention': np.zeros((1, ModelConstants.TRAFFIC_CONVENTION_LEN), dtype=np.float16),
+      'lateral_control_params': np.zeros((1, ModelConstants.LATERAL_CONTROL_PARAMS_LEN), dtype=np.float16),
+      'prev_desired_curv': np.zeros((1,(ModelConstants.HISTORY_BUFFER_LEN+1), ModelConstants.PREV_DESIRED_CURV_LEN), dtype=np.float16),
+      'features_buffer': np.zeros((1, ModelConstants.HISTORY_BUFFER_LEN,  ModelConstants.FEATURE_LEN), dtype=np.float16),
     }
 
     with open(METADATA_PATH, 'rb') as f:
@@ -77,9 +80,8 @@ def __init__(self, context: CLContext):
     self.output = np.zeros(net_output_size, dtype=np.float32)
     self.parser = Parser()
 
-    self.model = ModelRunner(MODEL_PATHS, self.output, Runtime.GPU, False, context)
-    for k,v in self.inputs.items():
-      self.model.addInput(k, v)
+    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()}
@@ -97,18 +99,15 @@ def run(self, buf: VisionBuf, wbuf: VisionBuf, transform: np.ndarray, transform_
 
     self.inputs['traffic_convention'][:] = inputs['traffic_convention']
     self.inputs['lateral_control_params'][:] = inputs['lateral_control_params']
-
-    new_img = self.frame.prepare(buf, transform.flatten(), self.model.getCLBuffer("input_imgs"))
-    self.model.setInputBuffer("input_imgs", new_img)
+    self.inputs['input_imgs'] = self.frame.prepare(buf, transform.flatten(), None).astype(np.float16).reshape(self.inputs['input_imgs'].shape)
     if wbuf is not None:
-      new_big_img = self.wide_frame.prepare(wbuf, transform_wide.flatten(), self.model.getCLBuffer("big_input_imgs"))
-      self.model.setInputBuffer("big_input_imgs", new_big_img)
-
+      self.inputs['big_input_imgs'] = self.wide_frame.prepare(wbuf, transform_wide.flatten(), None).astype(np.float16).reshape(self.inputs['input_imgs'].shape)
 
     if prepare_only:
       return None
 
-    self.model.execute()
+    self.tensor_inputs = {k: Tensor(v) for k, v in self.inputs.items()}
+    self.output = self.model_run(**self.tensor_inputs)['outputs'].numpy().flatten()
     outputs = self.parser.parse_outputs(self.slice_outputs(self.output))
 
     self.inputs['features_buffer'][:-ModelConstants.FEATURE_LEN] = self.inputs['features_buffer'][ModelConstants.FEATURE_LEN:]

selfdrive/modeld/runners/__init__.py

@@ -3,18 +3,18 @@
 from openpilot.selfdrive.modeld.runners.runmodel_pyx import RunModel, Runtime
 assert Runtime
 
-USE_THNEED = int(os.getenv('USE_THNEED', str(int(TICI))))
+USE_TINYGRAD = int(os.getenv('USE_TINYGRAD', str(int(TICI))))
 USE_SNPE = int(os.getenv('USE_SNPE', str(int(TICI))))
 
 class ModelRunner(RunModel):
-  THNEED = 'THNEED'
+  TINYGRAD = 'TINYGRAD'
   SNPE = 'SNPE'
   ONNX = 'ONNX'
 
   def __new__(cls, paths, *args, **kwargs):
-    if ModelRunner.THNEED in paths and USE_THNEED:
-      from openpilot.selfdrive.modeld.runners.thneedmodel_pyx import ThneedModel as Runner
-      runner_type = ModelRunner.THNEED
+    if ModelRunner.TINYGRAD in paths and USE_TINYGRAD:
+      from openpilot.selfdrive.modeld.runners.tinygradmodel import TinygradModel as Runner
+      runner_type = ModelRunner.TINYGRAD
     elif ModelRunner.SNPE in paths and USE_SNPE:
       from openpilot.selfdrive.modeld.runners.snpemodel_pyx import SNPEModel as Runner
       runner_type = ModelRunner.SNPE