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

@@ -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:

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)