Calibrate in tg

common/transformations/camera.py

@@ -177,3 +177,18 @@ def img_from_device(pt_device):
   pt_img = pt_view/pt_view[:, 2:3]
   return pt_img.reshape(input_shape)[:, :2]
 
+
+# Get venus stride buffer parameters based on resolution
+def get_nv12_info(width: int, height: int) -> tuple[int, int, int]:
+  if width == 1928 and height == 1208:
+    STRIDE = 2048
+    UV_OFFSET = 1216 * STRIDE
+    YUV_SIZE = 2346 * STRIDE
+    return YUV_SIZE, STRIDE, UV_OFFSET
+  elif width == 1344 and height == 760:
+    STRIDE = 1408
+    UV_OFFSET = 760 * STRIDE
+    YUV_SIZE = 2900 * STRIDE
+    return YUV_SIZE, STRIDE, UV_OFFSET
+  else:
+    raise NotImplementedError(f"Unsupported resolution for vipc: {width}x{height}")

selfdrive/modeld/SConscript

@@ -10,8 +10,6 @@ frameworks = []
 
 common_src = [
   "models/commonmodel.cc",
-  "transforms/loadyuv.cc",
-  "transforms/transform.cc",
 ]
 
 # OpenCL is a framework on Mac
@@ -20,15 +18,6 @@ if arch == "Darwin":
 else:
   libs += ['OpenCL']
 
-# Set path definitions
-for pathdef, fn in {'TRANSFORM': 'transforms/transform.cl', 'LOADYUV': 'transforms/loadyuv.cl'}.items():
-  for xenv in (lenv, lenvCython):
-    xenv['CXXFLAGS'].append(f'-D{pathdef}_PATH=\\"{File(fn).abspath}\\"')
-
-# Compile cython
-cython_libs = envCython["LIBS"] + libs
-commonmodel_lib = lenv.Library('commonmodel', common_src)
-lenvCython.Program('models/commonmodel_pyx.so', 'models/commonmodel_pyx.pyx', LIBS=[commonmodel_lib, *cython_libs], FRAMEWORKS=frameworks)
 tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=env.Dir("#").abspath) if 'pycache' not in x]
 
 # Get model metadata
@@ -38,22 +27,30 @@ for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
   cmd = f'python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
   lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files, cmd)
 
+# compile warp
+tg_flags = {
+    'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0',
+    'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}', # tinygrad calls brew which needs a $HOME in the env
+}.get(arch, 'DEV=CPU CPU_LLVM=1')
+image_flag = {
+     'larch64': 'IMAGE=2',
+}.get(arch, 'IMAGE=0')
+script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)]
+cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py '
+lenv.Command(fn + "warp_tinygrad.pkl", tinygrad_files + script_files, cmd)
+
 def tg_compile(flags, model_name):
   pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
   fn = File(f"models/{model_name}").abspath
   return lenv.Command(
     fn + "_tinygrad.pkl",
     [fn + ".onnx"] + tinygrad_files,
-    f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {fn}_tinygrad.pkl'
+    f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {fn}_tinygrad.pkl'
   )
 
 # Compile small models
 for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
-  flags = {
-    'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 IMAGE=2 JIT_BATCH_SIZE=0',
-    'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}', # tinygrad calls brew which needs a $HOME in the env
-  }.get(arch, 'DEV=CPU CPU_LLVM=1')
-  tg_compile(flags, model_name)
+  tg_compile(tg_flags, model_name)
 
 # Compile BIG model if USB GPU is available
 if "USBGPU" in os.environ:

selfdrive/modeld/compile_warp.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+import time
+import pickle
+import numpy as np
+from pathlib import Path
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import Context
+from tinygrad.device import Device
+from openpilot.common.transformations.camera import get_nv12_info
+
+
+WARP_PKL_PATH = Path(__file__).parent / 'models/warp_tinygrad.pkl'
+DM_WARP_PKL_PATH = Path(__file__).parent / 'models/dm_warp_tinygrad.pkl'
+
+MODEL_WIDTH = 512
+MODEL_HEIGHT = 256
+MODEL_FRAME_SIZE = MODEL_WIDTH * MODEL_HEIGHT * 3 // 2
+IMG_BUFFER_SHAPE = (30, 128, 256)
+W, H = 1928, 1208
+
+YUV_SIZE, STRIDE, UV_OFFSET = get_nv12_info(W, H)
+
+UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32)
+UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX)
+
+
+def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, ratio):
+  w_dst, h_dst = dst_shape
+  h_src, w_src = src_shape
+
+  x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst)
+  y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst)
+  ones = Tensor.ones_like(x)
+  dst_coords = x.reshape(1, -1).cat(y.reshape(1, -1)).cat(ones.reshape(1, -1))
+
+  src_coords = M_inv @ dst_coords
+  src_coords = src_coords / src_coords[2:3, :]
+
+  x_nn_clipped = Tensor.round(src_coords[0]).clip(0, w_src - 1).cast('int')
+  y_nn_clipped = Tensor.round(src_coords[1]).clip(0, h_src - 1).cast('int')
+  idx = y_nn_clipped * w_src + (y_nn_clipped * ratio).cast('int') * stride_pad + x_nn_clipped
+
+  sampled = src_flat[idx]
+  return sampled
+
+def frames_to_tensor(frames):
+  H = (frames.shape[0]*2)//3
+  W = frames.shape[1]
+  in_img1 = Tensor.cat(frames[0:H:2, 0::2],
+                        frames[1:H:2, 0::2],
+                        frames[0:H:2, 1::2],
+                        frames[1:H:2, 1::2],
+                        frames[H:H+H//4].reshape((H//2,W//2)),
+                        frames[H+H//4:H+H//2].reshape((H//2,W//2)), dim=0).reshape((6, H//2, W//2))
+  return in_img1
+
+def frame_prepare_tinygrad(input_frame, M_inv):
+  tg_scale = Tensor(UV_SCALE_MATRIX)
+  M_inv_uv = tg_scale @ M_inv @ Tensor(UV_SCALE_MATRIX_INV)
+  with Context(SPLIT_REDUCEOP=0):
+    y = warp_perspective_tinygrad(input_frame[:H*STRIDE],
+                                  M_inv, (MODEL_WIDTH, MODEL_HEIGHT),
+                                  (H, W), STRIDE - W, 1).realize()
+    u = warp_perspective_tinygrad(input_frame[UV_OFFSET:UV_OFFSET + (H//4)*STRIDE],
+                                  M_inv_uv, (MODEL_WIDTH//2, MODEL_HEIGHT//2),
+                                  (H//2, W//2), STRIDE - W, 0.5).realize()
+    v = warp_perspective_tinygrad(input_frame[UV_OFFSET + (H//4)*STRIDE:UV_OFFSET + (H//2)*STRIDE],
+                                  M_inv_uv, (MODEL_WIDTH//2, MODEL_HEIGHT//2),
+                                  (H//2, W//2), STRIDE - W, 0.5).realize()
+  yuv = y.cat(u).cat(v).reshape((MODEL_HEIGHT*3//2,MODEL_WIDTH))
+  tensor = frames_to_tensor(yuv)
+  return tensor
+
+def update_img_input_tinygrad(tensor, frame, M_inv):
+  M_inv = M_inv.to(Device.DEFAULT)
+  new_img = frame_prepare_tinygrad(frame, M_inv)
+  full_buffer = tensor[6:].cat(new_img, dim=0).contiguous()
+  return full_buffer, Tensor.cat(full_buffer[:6], full_buffer[-6:], dim=0).contiguous().reshape(1,12,MODEL_HEIGHT//2,MODEL_WIDTH//2)
+
+def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv,
+                              calib_big_img_buffer, new_big_img, M_inv_big):
+  calib_img_buffer, calib_img_pair = update_img_input_tinygrad(calib_img_buffer, new_img, M_inv)
+  calib_big_img_buffer, calib_big_img_pair = update_img_input_tinygrad(calib_big_img_buffer, new_big_img, M_inv_big)
+  return calib_img_buffer, calib_img_pair, calib_big_img_buffer, calib_big_img_pair
+
+def warp_perspective_numpy(src, M_inv, dst_shape, src_shape, stride_pad, ratio):
+    w_dst, h_dst = dst_shape
+    h_src, w_src = src_shape
+    xs, ys = np.meshgrid(np.arange(w_dst), np.arange(h_dst))
+
+    ones = np.ones_like(xs)
+    dst_hom = np.stack([xs, ys, ones], axis=0).reshape(3, -1)
+
+    src_hom = M_inv @ dst_hom
+    src_hom /= src_hom[2:3, :]
+
+    src_x = np.clip(np.round(src_hom[0, :]).astype(int), 0, w_src - 1)
+    src_y = np.clip(np.round(src_hom[1, :]).astype(int), 0, h_src - 1)
+    idx = src_y * w_src + (src_y * ratio).astype(np.int32) * stride_pad + src_x
+    return src[idx]
+
+
+def frames_to_tensor_np(frames):
+  H = (frames.shape[0]*2)//3
+  W = frames.shape[1]
+  p1 = frames[0:H:2, 0::2]
+  p2 = frames[1:H:2, 0::2]
+  p3 = frames[0:H:2, 1::2]
+  p4 = frames[1:H:2, 1::2]
+  p5 = frames[H:H+H//4].reshape((H//2, W//2))
+  p6 = frames[H+H//4:H+H//2].reshape((H//2, W//2))
+  return np.concatenate([p1, p2, p3, p4, p5, p6], axis=0)\
+           .reshape((6, H//2, W//2))
+
+def frame_prepare_np(input_frame, M_inv):
+  M_inv_uv = UV_SCALE_MATRIX @ M_inv @ UV_SCALE_MATRIX_INV
+  y  = warp_perspective_numpy(input_frame[:H*STRIDE],
+                                 M_inv, (MODEL_WIDTH, MODEL_HEIGHT), (H, W), STRIDE - W, 1)
+  u  = warp_perspective_numpy(input_frame[UV_OFFSET:UV_OFFSET + (H//4)*STRIDE],
+                                 M_inv_uv, (MODEL_WIDTH//2, MODEL_HEIGHT//2), (H//2, W//2), STRIDE - W, 0.5)
+  v  = warp_perspective_numpy(input_frame[UV_OFFSET + (H//4)*STRIDE:UV_OFFSET + (H//2)*STRIDE],
+                                 M_inv_uv, (MODEL_WIDTH//2, MODEL_HEIGHT//2), (H//2, W//2), STRIDE - W, 0.5)
+  yuv = np.concatenate([y, u, v]).reshape( MODEL_HEIGHT*3//2, MODEL_WIDTH)
+  return frames_to_tensor_np(yuv)
+
+def update_img_input_np(tensor, frame, M_inv):
+  tensor[:-6]  = tensor[6:]
+  tensor[-6:] = frame_prepare_np(frame, M_inv)
+  return tensor, np.concatenate([tensor[:6], tensor[-6:]], axis=0).reshape((1,12,MODEL_HEIGHT//2, MODEL_WIDTH//2))
+
+def update_both_imgs_np(calib_img_buffer, new_img, M_inv,
+                        calib_big_img_buffer, new_big_img, M_inv_big):
+  calib_img_buffer, calib_img_pair = update_img_input_np(calib_img_buffer, new_img, M_inv)
+  calib_big_img_buffer, calib_big_img_pair = update_img_input_np(calib_big_img_buffer, new_big_img, M_inv_big)
+  return calib_img_buffer, calib_img_pair, calib_big_img_buffer, calib_big_img_pair
+
+def run_and_save_pickle():
+  from tinygrad.engine.jit import TinyJit
+  from tinygrad.device import Device
+  update_img_jit = TinyJit(update_both_imgs_tinygrad, prune=True)
+
+  full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize()
+  big_full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize()
+  full_buffer_np = np.zeros(IMG_BUFFER_SHAPE, dtype=np.uint8)
+  big_full_buffer_np = np.zeros(IMG_BUFFER_SHAPE, dtype=np.uint8)
+
+  step_times = []
+  for _ in range(10):
+    new_frame_np = (32*np.random.randn(YUV_SIZE).astype(np.float32) + 128).clip(0,255).astype(np.uint8)
+    img_inputs = [full_buffer,
+                  Tensor.from_blob(new_frame_np.ctypes.data, (YUV_SIZE,), dtype='uint8').realize(),
+                  Tensor(Tensor.randn(3,3).mul(8).realize().numpy(), device='NPY')]
+    new_big_frame_np = (32*np.random.randn(YUV_SIZE).astype(np.float32) + 128).clip(0,255).astype(np.uint8)
+    big_img_inputs = [big_full_buffer,
+                      Tensor.from_blob(new_big_frame_np.ctypes.data, (YUV_SIZE,), dtype='uint8').realize(),
+                      Tensor(Tensor.randn(3,3).mul(8).realize().numpy(), device='NPY')]
+    inputs = img_inputs + big_img_inputs
+    Device.default.synchronize()
+    inputs_np = [x.numpy() for x in inputs]
+    inputs_np[0] = full_buffer_np
+    inputs_np[3] = big_full_buffer_np
+    st = time.perf_counter()
+    out = update_img_jit(*inputs)
+    full_buffer = out[0].contiguous().realize().clone()
+    big_full_buffer = out[2].contiguous().realize().clone()
+    mt = time.perf_counter()
+    Device.default.synchronize()
+    et = time.perf_counter()
+    step_times.append((et-st)*1e3)
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
+    out_np = update_both_imgs_np(*inputs_np)
+    full_buffer_np = out_np[0]
+    big_full_buffer_np = out_np[2]
+
+    # TODO REACTIVATE
+    #for a, b in zip(out_np, (x.numpy() for x in out), strict=True):
+    #  mismatch = np.abs(a - b) > 0
+    #  mismatch_percent = sum(mismatch.flatten()) / len(mismatch.flatten()) * 100
+    #  mismatch_percent_tol = 1e-2
+    #  assert mismatch_percent < mismatch_percent_tol, f"input mismatch percent {mismatch_percent} exceeds tolerance {mismatch_percent_tol}"
+
+  with open(WARP_PKL_PATH, "wb") as f:
+    pickle.dump(update_img_jit, f)
+
+  jit = pickle.load(open(WARP_PKL_PATH, "rb"))
+  # test function after loading
+  jit(*inputs)
+
+
+  def warp_dm(input_frame, M_inv):
+    M_inv = M_inv.to(Device.DEFAULT)
+    with Context(SPLIT_REDUCEOP=0):
+      result = warp_perspective_tinygrad(input_frame[:H*STRIDE], M_inv, (1440, 960), (H, W), STRIDE - W, 1).reshape(-1,960*1440)
+    return result
+  warp_dm_jit = TinyJit(warp_dm, prune=True)
+  step_times = []
+  for _ in range(10):
+    inputs = [Tensor.from_blob((32*Tensor.randn(YUV_SIZE,) + 128).cast(dtype='uint8').realize().numpy().ctypes.data, (YUV_SIZE,), dtype='uint8'),
+                  Tensor(Tensor.randn(3,3).mul(8).realize().numpy(), device='NPY')]
+
+    Device.default.synchronize()
+    st = time.perf_counter()
+    out = warp_dm_jit(*inputs)
+    mt = time.perf_counter()
+    Device.default.synchronize()
+    et = time.perf_counter()
+    step_times.append((et-st)*1e3)
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
+
+  with open(DM_WARP_PKL_PATH, "wb") as f:
+    pickle.dump(warp_dm_jit, f)
+
+if __name__ == "__main__":
+    run_and_save_pickle()

selfdrive/modeld/dmonitoringmodeld.py

@@ -3,7 +3,6 @@
 from openpilot.system.hardware import TICI
 os.environ['DEV'] = 'QCOM' if TICI else 'CPU'
 from tinygrad.tensor import Tensor
-from tinygrad.dtype import dtypes
 import time
 import pickle
 import numpy as np
@@ -12,19 +11,18 @@
 from cereal import messaging
 from cereal.messaging import PubMaster, SubMaster
 from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf
+from msgq.visionipc.visionipc_pyx import CLContext
 from openpilot.common.swaglog import cloudlog
 from openpilot.common.realtime import config_realtime_process
 from openpilot.common.transformations.model import dmonitoringmodel_intrinsics
-from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
-from openpilot.selfdrive.modeld.models.commonmodel_pyx import CLContext, MonitoringModelFrame
+from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye, get_nv12_info
 from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp
-from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address
 
 PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 MODEL_PKL_PATH = Path(__file__).parent / 'models/dmonitoring_model_tinygrad.pkl'
 METADATA_PATH = Path(__file__).parent / 'models/dmonitoring_model_metadata.pkl'
-
+DM_WARP_PKL_PATH = Path(__file__).parent / 'models/dm_warp_tinygrad.pkl'
 
 class ModelState:
   inputs: dict[str, np.ndarray]
@@ -36,28 +34,31 @@ def __init__(self, cl_ctx):
       self.input_shapes = model_metadata['input_shapes']
       self.output_slices = model_metadata['output_slices']
 
-    self.frame = MonitoringModelFrame(cl_ctx)
     self.numpy_inputs = {
       'calib': np.zeros(self.input_shapes['calib'], dtype=np.float32),
     }
 
+    self.warp_inputs_np = {'transform': np.zeros((3,3), dtype=np.float32)}
+    self.warp_inputs = {k: Tensor(v, device='NPY') for k,v in self.warp_inputs_np.items()}
+    self.frame_buf_params = None
     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)
 
+    with open(DM_WARP_PKL_PATH, "rb") as f:
+      self.image_warp = pickle.load(f)
+
   def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]:
     self.numpy_inputs['calib'][0,:] = calib
 
     t1 = time.perf_counter()
 
-    input_img_cl = self.frame.prepare(buf, transform.flatten())
-    if TICI:
-      # The imgs tensors are backed by opencl memory, only need init once
-      if 'input_img' not in self.tensor_inputs:
-        self.tensor_inputs['input_img'] = qcom_tensor_from_opencl_address(input_img_cl.mem_address, self.input_shapes['input_img'], dtype=dtypes.uint8)
-    else:
-      self.tensor_inputs['input_img'] = Tensor(self.frame.buffer_from_cl(input_img_cl).reshape(self.input_shapes['input_img']), dtype=dtypes.uint8).realize()
+    if self.frame_buf_params is None:
+      self.frame_buf_params = get_nv12_info(buf.width, buf.height)
+    self.warp_inputs['frame'] = Tensor.from_blob(buf.data.ctypes.data, (self.frame_buf_params[0],), dtype='uint8').realize()
 
+    self.warp_inputs_np['transform'][:] = transform[:]
+    self.tensor_inputs['input_img'] = self.image_warp(self.warp_inputs['frame'], self.warp_inputs['transform']).realize()
 
     output = self.model_run(**self.tensor_inputs).contiguous().realize().uop.base.buffer.numpy()