Repository: tinygrad/tinygrad
Pull request: fix(beam): GlobalCounters kernel count increasing when clearing l2
fixes #4595 by not updating stats when clearing l2 cache in beam search. Added a NOSTATS env var to prevent stats update with a context when some operations need to be ran without updating framework's stats.

Problem statement:
BEAM kernel count number is wrong
```
beam2  : 16 31 31 16 2 3 2 4 3 2        :   817.92 us  <  hc     : 4 31 31 32 4 3 3 4 2 2         :  1000.83 us
*** GPU        9 r_16_31_31_16_2_3_2_4_3_2              arg   3 mem  0.87 GB tm   1244.89us/     4.99ms (  113.83 GFLOPS,   24.03 GB/s)
   0.00s:                 from   1 ->   1 actions  512   32   32   32   12
   4.62s:      7940.00 us from  30 ->  30 actions  512    2   32   32   16   12
  13.97s:      3760.83 us from  59 ->  59 actions  512    2   32    8   16   12    4
  23.99s:      2318.33 us from  56 ->  56 actions  512    2   32   16   12    4    8
  35.39s:      1741.56 us from  47 ->  47 actions  512   32    4   16    2   12    4    2
  43.81s:      1741.56 us from  41 ->  41 actions  512   32    4   16    2   12    4    2
beam2  : 512 32 4 16 2 12 4 2           :  1778.44 us  <  hc     : 512 32 2 16 8 4 12             :  3710.63 us
*** GPU       16 r_512_32_4_16_2_12_4_2                 arg   3 mem  0.96 GB tm   1835.73us/     6.83ms (  274.18 GFLOPS,   49.42 GB/s)
```

Kernel number jumps from 9 to 16, it should only go up by one (like from 9 to 10)

Hints from issue discussion:
The problem is with this realize, will make the pr after my final exam (in 11mins).
<img width="1000" alt="image" src="https://github.com/tinygrad/tinygrad/assets/44068562/e77f442b-de5d-4dee-b2d3-68a3a94a13de">

Commit message: fix(test_search): added assert message

Patch to review:
diff --git a/test/test_search.py b/test/test_search.py
index baac49136e791..29b1f4f0f5a20 100644
--- a/test/test_search.py
+++ b/test/test_search.py
@@ -3,11 +3,11 @@
 from tinygrad.codegen.kernel import Opt, OptOps
 from tinygrad.codegen.linearizer import Linearizer
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.engine.search import time_linearizer, bufs_from_lin, actions
+from tinygrad.engine.search import time_linearizer, bufs_from_lin, actions, _ensure_buffer_alloc, _time_program
 from tinygrad.device import Device, Buffer
 from tinygrad.ops import LoadOps, BufferOps
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, GlobalCounters, to_function_name
 from tinygrad.engine.realize import capturing
 
 class TestTimeLinearizer(unittest.TestCase):
@@ -66,5 +66,22 @@ def test_get_linearizer_actions(self):
     if Opt(OptOps.GROUPTOP, 0, 0) in actions:
       assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUPTOP, axis=0, amt=3)]) == 0, "did not de-dup GROUPTOP"
 
+  def test_kernel_count(self):
+    """
+    Ensure that the kernel count is not incremented by _time_program when clearing l2
+    """
+    dev = Device[Device.DEFAULT]
+    assert dev.compiler is not None
+
+    si = [i for i in create_schedule([Tensor([1,2,3,4]).add(1).lazydata]) if i.ast[0].op not in LoadOps][0]
+    rawbufs = _ensure_buffer_alloc(bufs_from_lin(lin:=Linearizer(*si.ast)))
+    var_vals = {k:(k.max+k.min)//2 for k in lin.ast[0].vars()}
+    p = lin.to_program()
+
+    kernel_count = GlobalCounters.kernel_count
+    _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
+        max_global_size=65536, clear_l2=True, cnt=3, name=to_function_name(lin.name))
+    assert GlobalCounters.kernel_count == kernel_count, "kernel count was incremented by _time_program"
+
 if __name__ == '__main__':
   unittest.main()
diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py
index 2da25cda50004..0ca5e27cf4edf 100644
--- a/tinygrad/engine/realize.py
+++ b/tinygrad/engine/realize.py
@@ -182,7 +182,7 @@ def lower_schedule(schedule:List[ScheduleItem]) -> Generator[ExecItem, None, Non
 
 capturing: List = []  # put classes with an add method in here
 
-def run_schedule(schedule:List[ScheduleItem], var_vals:Optional[Dict[Variable, int]]=None):
+def run_schedule(schedule:List[ScheduleItem], var_vals:Optional[Dict[Variable, int]]=None, do_update_stats=True):
   for ei in lower_schedule(schedule):
     if len(capturing): capturing[0].add(ei)
-    ei.run(var_vals)
+    ei.run(var_vals, do_update_stats=do_update_stats)
diff --git a/tinygrad/engine/search.py b/tinygrad/engine/search.py
index d8950a16b1966..88ab17cc78423 100644
--- a/tinygrad/engine/search.py
+++ b/tinygrad/engine/search.py
@@ -44,7 +44,7 @@ def _time_program(p:Program, lib:bytes, var_vals, rawbufs, early_stop=None, max_
   tms = []
   for _ in range(cnt):
     if clear_l2:
-      with Context(DEBUG=0, BEAM=0, CACHECOLLECTING=0): Tensor.ones(1024,1024).contiguous().realize()
+      with Context(DEBUG=0, BEAM=0, CACHECOLLECTING=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
     tms.append(cast(float, car(rawbufs, var_vals, wait=True))*factor)
     if early_stop is not None and early_stop < tms[-1]: break
   return tms
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 71589842a51c1..4c45b25bbaaf9 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -160,9 +160,9 @@ def schedule(self, *lst:Tensor, seen:Optional[Set[LazyBuffer]]=None) -> List[Sch
     assert len(var_vals) == 0
     return schedule
 
-  def realize(self, *lst:Tensor) -> Tensor:
+  def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Trigger the computation needed to create these Tensor(s)."""
-    run_schedule(*self.schedule_with_vars(*lst))
+    run_schedule(*self.schedule_with_vars(*lst), do_update_stats=do_update_stats)
     return self
 
   def replace(self, x:Tensor) -> Tensor:


Code review comments:
diff_hunk: @@ -66,5 +66,22 @@ def test_get_linearizer_actions(self):
     if Opt(OptOps.GROUPTOP, 0, 0) in actions:
       assert len([x for x in lins if x.applied_opts[0] == Opt(OptOps.GROUPTOP, axis=0, amt=3)]) == 0, "did not de-dup GROUPTOP"
 
+  def test_kernel_count(self):
+    """
+    Ensure that the kernel count is not incremented by _time_program when clearing l2
+    """
+    dev = Device[Device.DEFAULT]
+    assert dev.compiler is not None
+
+    si = [i for i in create_schedule([Tensor([1,2,3,4]).add(1).lazydata]) if i.ast[0].op not in LoadOps][0]
+    rawbufs = _ensure_buffer_alloc(bufs_from_lin(lin:=Linearizer(*si.ast)))
+    var_vals = {k:(k.max+k.min)//2 for k in lin.ast[0].vars()}
+    p = lin.to_program()
+
+    kernel_count = GlobalCounters.kernel_count
+    _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
path: test/test_search.py
text: @user1:
prefer not to test internal functions directly (function name that start with `_`). can you use `time_linearizer` instead? also kernel count not changing may due to nothing is being run, i'd test the count increased by exactly one after running one kernel

@author:
I switched to testing `time_linearizer`. Though I don't understand the "running one kernel part". Is time_linearizer supposed to run kernels which count towards kcount ?

@user1:
it was 2 separate asks (1) don't call internal function in test, (2) test one kernel ran and the kernel count increased by 1 (not 7 or 10 in search).