BEAM bugfix, kernels dedup now (#5617)

* BEAM bugfix, kernels dedup now

* getenv is default

examples/handcode_opt.py

@@ -110,7 +110,7 @@ if __name__ == "__main__":
     # benchmark the programs
     choices = []
     for (lin, nm) in lins:
-      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
+      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
       ops = lin.to_program().op_estimate
       gflops = sym_infer(ops, {k:k.min for k in lin.ast.vars()})*1e-9/tm
       choices.append((tm, gflops, lin.linearize(), nm))

test/test_search.py

@@ -97,5 +97,16 @@ class TestBEAM(unittest.TestCase):
     tm = time_linearizer(best_lin, bufs, allow_test_size=False, cnt=2, disable_cache=True)
     assert tm
 
+  def test_beam_unnamed_kernels(self):
+    a = Tensor.rand(100)
+    b = Tensor.rand(100)
+    si = (a+b).schedule()[-1]
+    lin = Kernel(si.ast)
+    bufs = bufs_from_lin(lin)
+    # TODO: beam should have better instrumentation so we don't have to check this indirect thing
+    kcount = len(Kernel.kernel_cnt)
+    beam_search(lin, bufs, 3, disable_cache=True)
+    self.assertEqual(kcount, len(Kernel.kernel_cnt))
+
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/kernel.py

@@ -753,9 +753,9 @@ class Kernel:
     if getenv("GRAPHUOPS"): self.uops.graph()
     return self
 
-  def to_program(self) -> Program:
+  def to_program(self, name_override:Optional[str]=None) -> Program:
     self.linearize()
-    src = self.opts.render(name:=to_function_name(self.name), self.uops)
+    src = self.opts.render(name:=to_function_name(ansiname:=(name_override if name_override is not None else self.name)), self.uops)
     if getenv("RUN_PROCESS_REPLAY"):
       table_name = f"process_replay_{getenv('GITHUB_RUN_ID', 'HEAD')}"
       diskcache_put(table_name, id(self), (self.ast, self.opts, self.applied_opts, name, src, {k:v.value for k,v in ContextVar._cache.items()}))
@@ -765,4 +765,4 @@ class Kernel:
     mem_bytes = sum(max(x.arg.dtype.itemsize * x.arg.st.real_size() for x in group) for _, group in
       itertools.groupby([x for x in self.ast.lazyops if x.op in BufferOps and isinstance(x.arg, MemBuffer) and x.arg.idx >= 0],
                         key=lambda x: (x.op, x.arg.idx)))
-    return Program(self.name, src, self.opts.device, self.global_size, self.local_size, self.uops, ops * run_count, min(mem * run_count, mem_bytes))
+    return Program(ansiname, src, self.opts.device, self.global_size, self.local_size, self.uops, ops * run_count, min(mem * run_count, mem_bytes))

tinygrad/engine/search.py

@@ -48,7 +48,7 @@ def _time_program(p:Program, lib:bytes, var_vals, rawbufs, early_stop=None, max_
     if clear_l2:
       with Context(DEBUG=0, BEAM=0, CAPTURING=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
     tms.append(cast(float, car(input_bufs, var_vals, wait=True))*factor)
-    if early_stop is not None and early_stop < tms[-1]: break
+    if early_stop is not None and early_stop < min(tms): break
   return tms
 
 class TimeoutException(Exception): pass
@@ -61,7 +61,7 @@ def _try_compile_linearized_w_idx(x:Tuple[int,Kernel], compiler:Compiler) -> Tup
   try:
     x[1].linearize()
     if len(x[1].uops.uops) >= getenv("BEAM_UOPS_MAX", 3000) > 0: raise RuntimeError("too many uops")
-    p = x[1].to_program()
+    p = x[1].to_program(name_override="test")
     st = time.perf_counter()
     prog = compiler.compile(p.src)
     et = time.perf_counter() - st
@@ -116,10 +116,10 @@ def get_kernel_actions(lin:Kernel, include_0=True) -> Dict[int, Kernel]:
   return acted_lins
 
 beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
-def beam_search(lin:Kernel, rawbufs:List[Buffer], amt:int, allow_test_size=True) -> Kernel:
+def beam_search(lin:Kernel, rawbufs:List[Buffer], amt:int, allow_test_size=True, disable_cache=getenv("IGNORE_BEAM_CACHE")) -> Kernel:
   global beam_pool
   key = {"ast": lin.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": lin.opts.device, "suffix": lin.opts.suffix}
-  if not getenv("IGNORE_BEAM_CACHE") and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
+  if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
     ret = lin.copy()
     for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
     return ret