add name uop (#9149)

* add name uop, TODO: refactor renderer to use

* renderer uses name uop

* fix tests

* render

* ptx
This commit is contained in:
George Hotz
2025-02-18 15:26:58 +08:00
committed by GitHub
parent 2db8b4046a
commit a4dab3ec3f
18 changed files with 49 additions and 31 deletions
+3 -3
View File
@@ -40,7 +40,7 @@ class PythonProgram:
loop_ends: dict[int, int] = {}
while i < len(self.uops):
uop, dtype, idp, arg = self.uops[i]
void_ops = {Ops.STORE, Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF}
void_ops = {Ops.STORE, Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.NAME}
if uop is Ops.DEFINE_ACC: idp = [idp[0]]
inp = [ul[v] for v in idp if self.uops[v][0] not in void_ops]
dtp = [dl[v] for v in idp if self.uops[v][0] not in void_ops]
@@ -60,7 +60,7 @@ class PythonProgram:
loop_ends[idp[0]] = i
i = idp[0]
continue
if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF):
if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.NAME):
# in the python emulator, the warp is always in sync
i += 1
continue
@@ -196,7 +196,7 @@ class PythonRenderer(Renderer):
if getenv("EMULATE_INTEL"): self.device, self.suffix, self.tensor_cores = "INTEL", "INTEL", IntelRenderer.tensor_cores
if getenv("EMULATE_AMX"): self.device, self.tensor_cores = "CLANG", ClangRenderer.tensor_cores
def render(self, name:str, uops:list[UOp]) -> str:
def render(self, uops:list[UOp]) -> str:
lops = [(u.op, u.dtype, [uops.index(v) for v in u.src], u.arg) for u in uops]
return base64.b64encode(pickle.dumps(lops)).decode()