faster beam [run_process_replay] (#5718)

This commit is contained in:
George Hotz 2024-07-25 11:58:41 -07:00 committed by GitHub
parent eff7c5fd2c
commit 142b7fb22f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 2 additions and 8 deletions

View File

@ -112,9 +112,6 @@ class Kernel:
self.simplify_ones()
self.simplify_merge_adjacent()
# cache
self.applied_opts_cache: Optional[List[Opt]] = None
def copy(self):
ret = type(self).__new__(type(self))
@ -131,9 +128,6 @@ class Kernel:
self.applied_opts[:], self.group_for_reduces, self.upcasted, self.local_dims, self.dont_use_locals
ret.tensor_core, ret.tensor_core_opts, ret.bufs_for_tensor_core = self.tensor_core, self.tensor_core_opts, self.bufs_for_tensor_core
# uncached since linearize didn't run
ret.applied_opts_cache = None
return ret
@property

View File

@ -61,9 +61,9 @@ def _try_compile_linearized_w_idx(x:Tuple[int,Kernel], compiler:Compiler) -> Tup
# set timeout
signal.alarm(getenv("BEAM_TIMEOUT_SEC", 10))
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(name_override="test")
assert isinstance(p.uops, UOpGraph), "uop graph wasn't generated?"
if len(p.uops.uops) >= getenv("BEAM_UOPS_MAX", 3000) > 0: raise RuntimeError("too many uops")
st = time.perf_counter()
prog = compiler.compile(p.src)
et = time.perf_counter() - st