move renderer into options (#4514)

* move renderer into options

* fix tests

* renders are functions
This commit is contained in:
George Hotz
2024-05-10 10:01:51 -07:00
committed by GitHub
parent 7c630a9a53
commit 4eef1ee9bf
15 changed files with 36 additions and 38 deletions
+8 -6
View File
@@ -178,13 +178,15 @@ class PythonProgram:
i += 1
return time.perf_counter() - st
def PythonRenderer(name:str, uops:UOpGraph) -> str:
lops = [(u.uop, u.dtype, [uops.uops.index(v) for v in u.vin], u.arg) for u in uops]
return base64.b64encode(pickle.dumps(lops)).decode()
class PythonCompiler(Compiler):
compiler_opts = CompilerOptions("METAL", has_tensor_cores=True) if getenv("EMULATE_METAL") else \
(CompilerOptions("HSA", has_tensor_cores=True) if getenv("EMULATE_HSA") else \
(CompilerOptions("CUDA", has_tensor_cores=True) if getenv("EMULATE_CUDA") else CompilerOptions("PYTHON")))
def render(self, name:str, uops:UOpGraph) -> str:
lops = [(u.uop, u.dtype, [uops.uops.index(v) for v in u.vin], u.arg) for u in uops]
return base64.b64encode(pickle.dumps(lops)).decode()
compiler_opts = CompilerOptions("METAL", has_tensor_cores=True, renderer=PythonRenderer) if getenv("EMULATE_METAL") else \
(CompilerOptions("HSA", has_tensor_cores=True, renderer=PythonRenderer) if getenv("EMULATE_HSA") else \
(CompilerOptions("CUDA", has_tensor_cores=True, renderer=PythonRenderer) if getenv("EMULATE_CUDA") else \
CompilerOptions("PYTHON", renderer=PythonRenderer)))
def compile(self, src:str) -> bytes: return base64.b64decode(src)
class PythonAllocator(Allocator):