From 7af7467f22df9c2b232cd5b2d3f53e92cbf00819 Mon Sep 17 00:00:00 2001 From: Elias Wahl <82230675+Eliulm@users.noreply.github.com> Date: Mon, 18 Mar 2024 21:13:32 +0100 Subject: [PATCH] Fix: BEAM search with PTX fails (#3786) * Pass Uopgraph instead of list * Add search compile linearizer test * Rename * fix lint * Remove test. Add UOpGraph type --- tinygrad/device.py | 3 ++- tinygrad/features/search.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tinygrad/device.py b/tinygrad/device.py index 052de64353..f144f0e70b 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -7,6 +7,7 @@ from tinygrad.helpers import ansilen, DEBUG, getenv, colored, BEAM, NOOPT, all_i from tinygrad.helpers import prod from tinygrad.shape.symbolic import Variable, sym_infer, sint from tinygrad.ops import LazyOp, get_lazyop_info, GlobalCounters +from tinygrad.codegen.uops import UOpGraph from dataclasses import dataclass if TYPE_CHECKING: @@ -183,7 +184,7 @@ MallocAllocator = _MallocAllocator() class Compiler: linearizer_opts: ClassVar[LinearizerOptions] def __init__(self, cachekey:Optional[str]=None): self.cachekey = None if getenv("DISABLE_COMPILER_CACHE") else cachekey - def render(self, name:str, uops) -> str: raise NotImplementedError("need a render function") + def render(self, name:str, uops:UOpGraph) -> str: raise NotImplementedError("need a render function") def compile(self, src:str) -> bytes: raise NotImplementedError("need a compile function") def compile_cached(self, src:str) -> bytes: if self.cachekey is None or (lib := diskcache_get(self.cachekey, src)) is None: diff --git a/tinygrad/features/search.py b/tinygrad/features/search.py index 55795ceb13..1bd26fcc42 100644 --- a/tinygrad/features/search.py +++ b/tinygrad/features/search.py @@ -49,7 +49,7 @@ def _time_program(variables:List[Variable], outcount:int, rdev:Compiled, lib:byt def _compile_linearizer(compiler:Compiler, lin:Linearizer, name:Optional[str]=None) -> Tuple[bytes, Optional[List[int]], Optional[List[int]], List[Variable], int]: lin.linearize() - src = compiler.render(name if name is not None else to_function_name(lin.name), lin.uops.uops) # NOTE: these all have the same name for deduping + src = compiler.render(name if name is not None else to_function_name(lin.name), lin.uops) # NOTE: these all have the same name for deduping if DEBUG >= 5: print(src) return compiler.compile(src), lin.global_size, lin.local_size, lin.uops.vars(), len(lin.outbufs)