From a18c1f317897ce288ded89e2f13c8bdd21ea6fe5 Mon Sep 17 00:00:00 2001
From: George Hotz <george@comma.ai>
Date: Thu, 20 Oct 2022 13:46:52 -0700
Subject: [PATCH] zero out the inputs

---
 extra/thneed.py      | 5 +++--
 openpilot/compile.py | 6 +++++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/extra/thneed.py b/extra/thneed.py
index 985fe042a1..4f2bddd5d3 100644
--- a/extra/thneed.py
+++ b/extra/thneed.py
@@ -111,10 +111,10 @@ class Thneed:
       ptr = nptr
   
     # populate the cl_cache
-    for k in jdat['kernels']:
+    for i,k in enumerate(jdat['kernels']):
       kernel = prgs[k['name']]
       aaa = []
-      for i,(a,sz) in enumerate(zip(k['args'], k['args_size'])):
+      for j,(a,sz) in enumerate(zip(k['args'], k['args_size'])):
         if len(a) == 0:
           aa = cl.LocalMemory(sz)
         elif len(a) == 4:
@@ -124,6 +124,7 @@ class Thneed:
           a = a.encode('latin_1')
           aa = np.uint16(struct.unpack("H", a)[0])
         elif len(a) == 8:
+          #print(i,j,struct.unpack("Q", a.encode('latin_1'))[0])
           aa = bufs[a]
         aaa.append(aa)
       self.cl_cache.append((kernel, [k['global_work_size'], k['local_work_size'], *aaa]))
diff --git a/openpilot/compile.py b/openpilot/compile.py
index 8eadf0098f..45f0800300 100644
--- a/openpilot/compile.py
+++ b/openpilot/compile.py
@@ -150,6 +150,10 @@ def compile(dat, output_fn):
         # but the differences are in unused stuff
         print("THNEED ERROR")
         for i,(a,b) in enumerate(zip(t.cl_cache, nt.cl_cache)):
+          for j,(c,d) in enumerate(zip(a[1][2:], b[1][2:])):
+            if type(c) == cl.Buffer and c not in t.buffers_to_save:
+              CL.enqueue_copy(c, b'\x00'*c.size)
+              CL.enqueue_copy(d, b'\x00'*d.size)
           # rerun
           a[0].clprg(CL().cl_queue, *a[1])
           CL().cl_queue.finish()
@@ -167,7 +171,7 @@ def compile(dat, output_fn):
               if not (cc == dd).all():
                 print(f"mismatch in layer {i} arg {j}")
                 np.testing.assert_allclose(cc, dd)
-        assert False
+        assert False, "all layers match?"
       print("thneed self-test passed!")
     except ModuleNotFoundError:
       pass