diff --git a/examples/handcode_opt.py b/examples/handcode_opt.py index d227ae89ce..c90ede2b93 100644 --- a/examples/handcode_opt.py +++ b/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)) diff --git a/test/test_search.py b/test/test_search.py index bb7837fff6..70b10464cc 100644 --- a/test/test_search.py +++ b/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() diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index f4d2444fe9..f8f39b32be 100644 --- a/tinygrad/codegen/kernel.py +++ b/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)) diff --git a/tinygrad/engine/search.py b/tinygrad/engine/search.py index 2bd08d5017..0d15465e51 100644 --- a/tinygrad/engine/search.py +++ b/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