github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

fix: don't put kernels into cache when optimizing (#1409)

Browse Source
This commit is contained in:
wozeparrot
2023-08-02 18:17:16 -04:00
committed by GitHub
parent 9ffa1eb7e2
commit a367f71fea
2 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tinygrad/codegen/optimizer.py
View File

@@ -30,9 +30,9 @@ def kernel_optimize_search(k:Linearizer, create_k:Callable[[], Linearizer], runt
k.process()
apply_opt(k, x)
prg = k.codegen().build(runtime)
first_tm = prg.exec(k.bufs, force_wait=True)
first_tm = prg.exec(k.bufs, force_wait=True, optimizing=True)
if baseline*5 < first_tm*1000: return first_tm*1000 # very slow
tm = min([first_tm]+[prg.exec(k.bufs, force_wait=True) for _ in range(2)])*1000
tm = min([first_tm]+[prg.exec(k.bufs, force_wait=True, optimizing=True) for _ in range(2)])*1000
return tm
except Exception:
if DEBUG >= 3:
@@ -76,7 +76,7 @@ def kernel_optimize(k:Linearizer, create_k:Callable[[], Linearizer], runtime):
k = create_k()
hand_coded_optimizations(k)
prg = k.codegen().build(runtime)
return min([prg.exec(k.bufs, force_wait=True) for _ in range(5)])*1000
return min([prg.exec(k.bufs, force_wait=True, optimizing=True) for _ in range(5)])*1000
choice = kernel_optimize_search(k, create_k, runtime, get_baseline())
if global_db is not None:
global_db[skey] = choice

4
tinygrad/ops.py
View File

@@ -133,9 +133,9 @@ class ASTRunner:
self.clprg = runtime(self.name, self.prg, **self.runtime_args)
return self
def exec(self, bufs, force_wait=False) -> Optional[float]:
def exec(self, bufs, force_wait=False, optimizing=False) -> Optional[float]:
rawbufs = dedup([x.realized for x in bufs if buf_is_kernel_arg(x)])
if GlobalCounters.cache is not None: GlobalCounters.cache.append((self, rawbufs))
if GlobalCounters.cache is not None and not optimizing: GlobalCounters.cache.append((self, rawbufs))
return self(rawbufs, force_wait=force_wait)
def __call__(self, rawbufs:List[RawBuffer], jit=False, force_wait=False) -> Optional[float]: