IQ.Pilot Release Commit @ 763bad7

iqpilot/selfdrive/iqmodeld/tests/test_egpu_policy.py

@@ -0,0 +1,82 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+import os
+
+import numpy as np
+
+os.environ["DEV"] = "CPU"
+
+from iqpilot.selfdrive.iqmodeld.egpu_policy import PolicyRunner, make_run_policy, packed_layout, queue_shapes
+from iqpilot.selfdrive.iqmodeld.temporal_state import TemporalInputState
+
+SPEC = {
+  "img": ((1, 12, 8, 16), "uint8"),
+  "big_img": ((1, 12, 8, 16), "uint8"),
+  "desire_pulse": ((1, 25, 8), "float32"),
+  "traffic_convention": ((1, 2), "float32"),
+  "action_t": ((1, 2), "float32"),
+  "features_buffer": ((1, 24, 512), "float32"),
+}
+FS = 4
+OUT_LEN = 2580
+HIDDEN = slice(1064, 1576)
+
+
+def _pack(inputs):
+  from tinygrad.tensor import Tensor
+  parts = [inputs[k].cast("float32").reshape(-1) for k in ("img", "big_img", "features_buffer", "desire_pulse", "traffic_convention", "action_t")]
+  flat = Tensor.cat(*parts)
+  hidden = (flat[:512] * 0.001).reshape(1, 512)
+  return flat, hidden
+
+
+def _fake_model(inputs):
+  from tinygrad.tensor import Tensor
+  flat, hidden = _pack(inputs)
+  n = flat.shape[0]
+  head = flat[:min(n, HIDDEN.start)]
+  out = Tensor.cat(head.pad((0, HIDDEN.start - head.shape[0])), hidden.reshape(-1), Tensor.zeros(OUT_LEN - HIDDEN.stop, device="CPU"))
+  return {"outputs": out.reshape(1, -1)}
+
+
+class _Reference:
+  def __init__(self):
+    self.state = TemporalInputState(FS, SPEC)
+
+  def run(self, warped, desire, traffic, action_t):
+    inputs = self.state.push_and_materialize(warped, desire, traffic, action_t)
+    from tinygrad.tensor import Tensor
+    t = {k: Tensor(np.ascontiguousarray(v), device="CPU") for k, v in inputs.items()}
+    out = _fake_model(t)["outputs"].numpy().reshape(-1)
+    self.state.note_hidden_state(out, HIDDEN)
+    return out
+
+
+def test_policy_queues_match_temporal_state():
+  from tinygrad.engine.jit import TinyJit
+  jit = TinyJit(make_run_policy(_fake_model, SPEC, FS, "CPU"), prune=True)
+  runner = PolicyRunner(jit, SPEC, FS, HIDDEN, "CPU")
+  ref = _Reference()
+  rng = np.random.default_rng(3)
+  desire = np.zeros(8, dtype=np.float32)
+  for i in range(14):
+    warped = rng.integers(0, 256, (2, 6, 8, 16), dtype=np.int64).astype(np.uint8)
+    if i in (2, 3, 9):
+      desire[:] = 0
+      desire[1 + (i % 3)] = 1
+    elif i == 5:
+      desire[:] = 0
+    traffic = np.array([1.0, 0.0], dtype=np.float32) if i % 2 else np.array([0.0, 1.0], dtype=np.float32)
+    action_t = np.array([0.1 * i, 0.2], dtype=np.float32)
+    got = runner.run(warped, desire, traffic, action_t)
+    want = ref.run(warped, desire, traffic, action_t)
+    np.testing.assert_array_equal(got, want, err_msg=f"frame {i}")
+
+
+def test_layouts():
+  shapes, sizes = packed_layout(SPEC)
+  assert list(shapes) == ["desire", "traffic_convention", "action_t", "prev_feat"]
+  assert sum(sizes) == 8 + 2 + 2 + 512
+  q = queue_shapes(SPEC, FS)
+  assert q["img_q"][0] == (5, 6, 8, 16) and q["feat_q"][0] == (96, 1, 512) and q["desire_q"][0] == (100, 1, 8)