From 2d0bfdf01ce2f3460aff2a62df1189a0254b19c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Szymon=20O=C5=BC=C3=B3g?= <58388001+SzymonOzog@users.noreply.github.com> Date: Sun, 24 Mar 2024 22:54:45 +0100 Subject: [PATCH] ptx cleanup (#3893) --- tinygrad/renderer/assembly.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tinygrad/renderer/assembly.py b/tinygrad/renderer/assembly.py index c52689dd72..86f576eead 100644 --- a/tinygrad/renderer/assembly.py +++ b/tinygrad/renderer/assembly.py @@ -66,10 +66,8 @@ def uops_to_asm(lang:AssemblyLanguage, function_name:str, uops:UOpGraph) -> str: def ptr_ar(root): root.arg = '.shared' if root.vin[0].uop == UOps.DEFINE_LOCAL else '.global' # move this to the argL - if root.vin[0].dtype.itemsize > 1: - val = uops.add(UOps.CONST, dtypes.int, tuple(), arg=root.vin[0].dtype.itemsize, insert_before=uops.uops.index(root)) - ptr = uops.add(UOps.ALU, dtypes.int, (root.vin[1], val), arg=BinaryOps.MUL, insert_before=uops.uops.index(root)) - else: ptr = root.vin[1] + val = uops.add(UOps.CONST, dtypes.int, tuple(), arg=root.vin[0].dtype.itemsize, insert_before=uops.uops.index(root)) + ptr = uops.add(UOps.ALU, dtypes.int, (root.vin[1], val), arg=BinaryOps.MUL, insert_before=uops.uops.index(root)) if ptr.uop == UOps.CONST: root.vin = (root.vin[0], ptr) + root.vin[2:] else: zero = uops.add(UOps.CONST, dtypes.int, tuple(), arg=0, cachable=False, insert_before=uops.uops.index(root)) @@ -98,6 +96,7 @@ def uops_to_asm(lang:AssemblyLanguage, function_name:str, uops:UOpGraph) -> str: if o in u.vin and u is not n: u.vin = tuple(n if x == o else x for x in u.vin) if rew := matcher.rewrite(u): replace[u] = rew + uops.remove_childless(set(x for x in uops if x.uop in {UOps.DEFINE_GLOBAL, UOps.PHI, UOps.ENDIF, UOps.ENDLOOP, UOps.STORE})) def kk(*s: str): kernel.append("\n".join(s))