From be53ef4f0a5a817a4d4335dfaa49f0cef93afb30 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 27 Jun 2025 11:09:25 -0700 Subject: [PATCH] rename DEFINE_ACC -> DEFINE_REG (#11006) * rename DEFINE_ACC -> DEFINE_REG * add CMPEQ to groupops --- extra/assembly/assembly.py | 2 +- extra/backends/triton.py | 2 +- test/test_linearizer.py | 12 ++++++------ test/test_quantize_onnx.py | 2 +- tinygrad/codegen/devectorizer.py | 4 ++-- tinygrad/codegen/expander.py | 2 +- tinygrad/codegen/linearize.py | 4 ++-- tinygrad/renderer/cstyle.py | 4 ++-- tinygrad/renderer/llvmir.py | 2 +- tinygrad/renderer/ptx.py | 6 +++--- tinygrad/runtime/ops_dsp.py | 2 +- tinygrad/runtime/ops_python.py | 4 ++-- tinygrad/uop/__init__.py | 14 ++++++++------ tinygrad/uop/spec.py | 4 ++-- tinygrad/uop/symbolic.py | 2 +- tinygrad/viz/serve.py | 2 +- 16 files changed, 35 insertions(+), 33 deletions(-) diff --git a/extra/assembly/assembly.py b/extra/assembly/assembly.py index 688f2b9aac..ab7334b57e 100644 --- a/extra/assembly/assembly.py +++ b/extra/assembly/assembly.py @@ -156,7 +156,7 @@ def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]): lang.ins.append(AssemblyInstruction(Ops.ALU, out, [tmp], args)) else: lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[x] for x in vin], args)) - elif uop == Ops.DEFINE_ACC: + elif uop == Ops.DEFINE_REG: reg = lang.newreg(u, dtype=dtype) lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], args)) elif uop == Ops.SPECIAL: diff --git a/extra/backends/triton.py b/extra/backends/triton.py index f653341623..154cbc405e 100644 --- a/extra/backends/triton.py +++ b/extra/backends/triton.py @@ -88,7 +88,7 @@ def uops_to_triton(function_name:str, uops:List[UOp]): assert dtype is not None if len(vin) == 2: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.load({r[vin[0]]} + { fill_dims_for_idx(r[vin[1]], dims)}, mask = {render_valid(valid)})', dtype)}") else: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.where({r[vin[2]]}, tl.load({r[vin[0]]}+{fill_dims_for_idx(r[vin[1]],dims)} , mask={render_valid(valid+[r[vin[2]]])}), 0.0)', dtype)}") - elif uop == Ops.DEFINE_ACC: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}") + elif uop == Ops.DEFINE_REG: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}") elif uop == Ops.CONST: r[u] = define_scalar([], dtype, args) elif uop == Ops.ASSIGN: kk(f"{r[vin[0]]} = {r[vin[1]].replace('//', '/')}") diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 8fa0f51310..325b7d168a 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -109,7 +109,7 @@ class TestLinearizer(unittest.TestCase): def _test_no_nested_ranges(self, lins, skip=None): for l in lins: - range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_ACC]) + range_in_acc = flatten([[x for x in u.src if x.op is Ops.RANGE] for u in l.uops if u.op is Ops.DEFINE_REG]) ranges = [u.op for u in l.uops if (u.op is Ops.RANGE and u in range_in_acc) or (u.op is Ops.ENDRANGE and u.src[0] in range_in_acc)] for i,u in enumerate(ranges): if skip and i in skip: continue @@ -255,7 +255,7 @@ class TestLinearizer(unittest.TestCase): k.upcast() k.upcast() k.linearize() - accs = [u for u in k.uops if u.op is Ops.DEFINE_ACC] + accs = [u for u in k.uops if u.op is Ops.DEFINE_REG] stores = [u for u in k.uops if u.op is Ops.STORE] assert len(accs) == 0 # it's removed now assert len(stores) == 1 @@ -310,7 +310,7 @@ class TestLinearizer(unittest.TestCase): realized_ast = a.schedule()[-1].ast realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple())) program = get_program(realized_ast, Device[Device.DEFAULT].renderer) - local = [uop for uop in program.uops if uop.op is Ops.DEFINE_ACC] + local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG] assert local[0].dtype == acc_dtype def test_arg_acc_dtype(self): @@ -318,7 +318,7 @@ class TestLinearizer(unittest.TestCase): realized_ast = c.schedule()[-1].ast realized_ast = realized_ast.replace(arg=KernelInfo(opts_to_apply=tuple())) program = get_program(realized_ast, Device[Device.DEFAULT].renderer) - local = [uop for uop in program.uops if uop.op is Ops.DEFINE_ACC] + local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG] assert local[0].dtype == expected_dtype tests = ( @@ -993,7 +993,7 @@ class TestFloat4(unittest.TestCase): ]: ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts))) program = get_program(ast, Device[Device.DEFAULT].renderer) - count = len([uop for uop in program.uops if uop.op is Ops.DEFINE_ACC and uop.dtype == dtypes.float.vec(4)]) + count = len([uop for uop in program.uops if uop.op is Ops.DEFINE_REG and uop.dtype == dtypes.float.vec(4)]) assert count == expected, f"{count=}, {expected=}" @unittest.skip("this doesn't happen anymore") @@ -1015,7 +1015,7 @@ class TestFloat4(unittest.TestCase): ]: ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts))) program = get_program(ast, Device[Device.DEFAULT].renderer) - count = len([uop for uop in program.uops if uop.op is Ops.DEFINE_ACC and uop.dtype == dtypes.float.vec(2)]) + count = len([uop for uop in program.uops if uop.op is Ops.DEFINE_REG and uop.dtype == dtypes.float.vec(2)]) assert count == expected, f"{count=}, {expected=}" class TestHandCodedOpts(unittest.TestCase): diff --git a/test/test_quantize_onnx.py b/test/test_quantize_onnx.py index 7b1dd2747e..7da8c769c6 100644 --- a/test/test_quantize_onnx.py +++ b/test/test_quantize_onnx.py @@ -78,7 +78,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase): with Context(DONT_REALIZE_EXPAND=1, QUANTIZE=1): sched = run_onnx({"input":inp})["output"].schedule() ei = lower_schedule_item(sched[-2]) - daccs = [u for u in ei.prg.p.uops if u.op is Ops.DEFINE_ACC] + daccs = [u for u in ei.prg.p.uops if u.op is Ops.DEFINE_REG] assert all(u.dtype.scalar() is dtypes.int for u in daccs) @unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP") diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index fb8eb9c225..f73d7057e1 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -290,7 +290,7 @@ devectorize = PatternMatcher([ # no ALU on vectorized dtypes (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.ASSIGN), name="alu"), no_vectorized_alu), (UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma), - (UPat(Ops.DEFINE_ACC, name="acc"), no_vectorized_acc), + (UPat(Ops.DEFINE_REG, name="acc"), no_vectorized_acc), ]) pm_render = PatternMatcher([ @@ -329,7 +329,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp): assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}" # if we have a range if len(reduce_range) != 0: - acc = UOp(Ops.DEFINE_ACC, red.dtype, (red.const_like(identity_element(red.arg, red.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,)) + acc = UOp(Ops.DEFINE_REG, red.dtype, (red.const_like(identity_element(red.arg, red.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,)) lst = [acc] + lst # put acc as the first element ctx.acc_num += 1 ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst) diff --git a/tinygrad/codegen/expander.py b/tinygrad/codegen/expander.py index 17a5faa412..c690c81245 100644 --- a/tinygrad/codegen/expander.py +++ b/tinygrad/codegen/expander.py @@ -87,7 +87,7 @@ expander = PatternMatcher([ Ops.VECTORIZE, Ops.IF, Ops.REDUCE), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand), (UPat(Ops.CONTRACT, name="con"), do_contract), # vectorize DEFINE_ACC - (UPat(Ops.VECTORIZE, src=UPat(Ops.DEFINE_ACC, name="acc"), name="v"), + (UPat(Ops.VECTORIZE, src=UPat(Ops.DEFINE_REG, name="acc"), name="v"), lambda acc,v: acc.replace(dtype=v.dtype, src=(acc.src[0].broadcast(v.dtype.count),)+acc.src[1:])), # BARRIERs aren't actually expanded (UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)), diff --git a/tinygrad/codegen/linearize.py b/tinygrad/codegen/linearize.py index d205299e54..3736be967b 100644 --- a/tinygrad/codegen/linearize.py +++ b/tinygrad/codegen/linearize.py @@ -99,7 +99,7 @@ class BlockContext: ctx.child_ctxs[u] = tuple([y for y in store_context if y not in idx_context and y.op is Ops.RANGE]) else: ctx.child_ctxs[u] = () elif u.op is Ops.ASSIGN: - assert u.src[0].op is Ops.DEFINE_ACC + assert u.src[0].op is Ops.DEFINE_REG ctx.child_ctxs[u] = tuple([y for y in ctx.last_ctx(u.src[1]) if y not in u.src[0].src[1:]]) return ctx @@ -216,7 +216,7 @@ def remove_blockend(x:UOp): parent_block = parent_blocks[0] assert len(parent_blocks) == parent_block.arg.cnt # range needs DEFINE_ACC to be before the range (never in DEFINE_ACC for if) - early_ops, late_ops = partition(x.arg.lst, lambda y: y.op is Ops.DEFINE_ACC and x.arg.end in y.src) + early_ops, late_ops = partition(x.arg.lst, lambda y: y.op is Ops.DEFINE_REG and x.arg.end in y.src) # NOTE: we have to add a barrier at the start if barrier is used in the range if x.op is Ops.BLOCKEND and any(y.op is Ops.BARRIER for y in late_ops) and late_ops[-1].op is Ops.ENDRANGE: late_ops = [UOp(Ops.BARRIER)] + late_ops diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index f4a9d0024c..3a15047f98 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -8,7 +8,7 @@ from tinygrad.renderer import Renderer, TensorCore from tinygrad.codegen.devectorizer import no_vectorized_alu base_rewrite = PatternMatcher([ - (UPat(Ops.DEFINE_ACC, name="x"), lambda ctx,x: ctx[x.src[0]]), + (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: ctx[x.src[0]]), (UPat(Ops.ASSIGN, name="x"), lambda ctx,x: f"{ctx[x.src[0]]} = {ctx[x.src[1]]};"), (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"), (UPat((Ops.ENDIF, Ops.ENDRANGE)), lambda ctx: "}"), @@ -154,7 +154,7 @@ class CStyleLanguage(Renderer): else: prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const", Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.NOOP: "precast", - Ops.INDEX: "bidx", Ops.DEFINE_ACC: "acc", Ops.LOAD: "val"}.get(u.op, "alu") + Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu") r[u] = f"{prefix}{c[prefix]}" l = cast(str, self.string_rewrite.rewrite(u, ctx=self)) diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index d882ed7ae2..32fd79854b 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -179,7 +179,7 @@ class LLVMRenderer(Renderer): local_args.append(f"@{r[u][1:]} = internal unnamed_addr addrspace(3) global [{u.dtype.size} x {ldt(u.dtype)}] undef, align 16") kernel.append(f" {r[u]} = addrspacecast [{u.dtype.size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{u.dtype.size} x {ldt(u.dtype)}]*") elif u.op is Ops.ASSIGN: pass # assign is already handled by the first pass - elif u.op is Ops.DEFINE_ACC: r[u] = r[u.src[0]] # a define acc can be used and never be assigned to + elif u.op is Ops.DEFINE_REG: r[u] = r[u.src[0]] # a define acc can be used and never be assigned to elif u.op is Ops.CONST: r[u] = lconst(u.arg, u.dtype) elif u.op is Ops.CAST and (ldt(u.dtype) == ldt(u.src[0].dtype) or isinstance(u.dtype, PtrDType)): r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index 1b15708c84..bf730b2e77 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -108,9 +108,9 @@ string_rewrite = PatternMatcher([ (UPat(Ops.LOAD, name="x", src=(UPat.var('loc'),), allow_any_len=True), lambda ctx, x, loc: f"ld.{mem_type(x)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \ if x.dtype.count > 1 else f"ld.{mem_type(x)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"), - (UPat(Ops.DEFINE_ACC, name="x", src=(UPat.cvar("pred", dtype=dtypes.bool),), allow_any_len=True), lambda ctx, x, pred: [ + (UPat(Ops.DEFINE_REG, name="x", src=(UPat.cvar("pred", dtype=dtypes.bool),), allow_any_len=True), lambda ctx, x, pred: [ f"setp.ne.s16 {ctx.r[pred]}, {render_val(pred.arg, pred.dtype)}, 0;", f"mov.pred {ctx.r[x]}, {ctx.r[pred]};"]), - (UPat(Ops.DEFINE_ACC, name="x", src=(UPat.cvar("pred"),), allow_any_len=True), + (UPat(Ops.DEFINE_REG, name="x", src=(UPat.cvar("pred"),), allow_any_len=True), lambda ctx, x, pred: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(pred.arg, x.dtype)};"), (UPat(Ops.RANGE, name="x"), lambda ctx, x: [f"mov.u32 {ctx.r[x]}, 0;", "LOOP_" + f"{ctx.r[x][1:]}:"]), (UPat(Ops.ASSIGN, name="x", dtype=dtypes.bool), lambda ctx, x: [f"mov.pred {ctx.r[x.src[0]]}, {ctx.r[x.src[1]]};"]), @@ -201,7 +201,7 @@ class PTXRenderer(Renderer): [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.arg[3].itemsize)]] r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) for _ in range(u.dtype.count)] prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.ENDRANGE: ("pred", "pred"), Ops.RANGE: ("ridx", None), - Ops.DEFINE_ACC: ("acc", None), Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL:("local",self.types[dtypes.ulong]), + Ops.DEFINE_REG: ("acc", None), Ops.DEFINE_VAR: ("dat", None), Ops.CONST: ("const", None), Ops.DEFINE_LOCAL:("local",self.types[dtypes.ulong]), Ops.DEFINE_GLOBAL: ("dat", self.types[dtypes.ulong]), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None)) if prefix: r[u] = ssa(prefix, u, dtype) diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index c62fba07bf..ff4ed21ecf 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -23,7 +23,7 @@ dsp_pm_late = PatternMatcher([ (UPat.var("x")+UPat(Ops.VECTORIZE,src=UPat.var("y")), lambda x,y: x+UOp(Ops.CUSTOMI,x.dtype,(y,),arg="{0}") if x.op is not Ops.CUSTOMI else None), (UPat.var("x")*UPat(Ops.VECTORIZE,src=UPat.var("y")), lambda x,y: x*UOp(Ops.CUSTOMI,x.dtype,(y,),arg="{0}") if x.op is not Ops.CUSTOMI else None), (UPat.var("x")//UPat(Ops.VECTORIZE,src=UPat.var("y")), lambda x,y: x//UOp(Ops.CUSTOMI,x.dtype,(y,),arg="{0}") if x.op is not Ops.CUSTOMI else None), - (UPat(Ops.DEFINE_ACC, src=(UPat(Ops.VECTORIZE, src=UPat(Ops.CONST, arg=0)),), dtype=dtypes.uchar.vec(128), name="d", allow_any_len=True), + (UPat(Ops.DEFINE_REG, src=(UPat(Ops.VECTORIZE, src=UPat(Ops.CONST, arg=0)),), dtype=dtypes.uchar.vec(128), name="d", allow_any_len=True), lambda d: d.replace(src=(UOp(Ops.CUSTOMI, d.dtype, arg="__builtin_HEXAGON_V6_vd0_128B()"),)+d.src[1:])), ]) diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 0cdc598507..147d315b9a 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -41,7 +41,7 @@ class PythonProgram: while i < len(self.uops): uop, dtype, idp, arg = self.uops[i] void_ops = {Ops.STORE, Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK} - if uop is Ops.DEFINE_ACC: idp = [idp[0]] + if uop is Ops.DEFINE_REG: 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] if getenv("TRACE"): print(i, uop, dtype, arg, inp, dtp) @@ -77,7 +77,7 @@ class PythonProgram: if arg[0][0] == 'g': ul[i] = [idxs[2-int(arg[0][-1])]] * warp_size elif arg[0][0] == 'l': ul[i] = [x[2-int(arg[0][-1])] for x in warp] elif uop is Ops.CONST: ul[i] = [arg] * warp_size - elif uop is Ops.DEFINE_ACC: + elif uop is Ops.DEFINE_REG: ul[i] = [[inp[0][0][0]] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [inp[0][0]] * warp_size elif uop is Ops.INDEX: ret:list = [] diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 3f69d702dc..ace8fee56b 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -31,7 +31,7 @@ class Ops(FastEnum): VALID = auto() # TODO: unify these ops into the levels of the memory hierarchy. depends on ASSIGN is STORE - DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_ACC = auto() # noqa: E702 + DEFINE_GLOBAL = auto(); DEFINE_LOCAL = auto(); DEFINE_REG = auto() # noqa: E702 # this is for symbolic shapes DEFINE_VAR = auto(); BIND = auto() # noqa: E702 @@ -59,8 +59,10 @@ class Ops(FastEnum): INDEX = auto() # BinaryOps - ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto(); CMPLT = auto(); CMPNE = auto() # noqa: E702 - EQ = auto(); XOR = auto(); OR = auto(); AND = auto(); THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto() # noqa: E702 + ADD = auto(); MUL = auto(); SHL = auto(); SHR = auto(); IDIV = auto(); MAX = auto(); MOD = auto() # noqa: E702 + CMPLT = auto(); CMPNE = auto(); CMPEQ = auto() # noqa: E702 + XOR = auto(); OR = auto(); AND = auto() # noqa: E702 + THREEFRY = auto(); SUB = auto(); FDIV = auto(); POW = auto() # noqa: E702 # TernaryOps WHERE = auto(); MULACC = auto() # noqa: E702 @@ -76,8 +78,8 @@ class Ops(FastEnum): class GroupOp: Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIP, Ops.NEG} - Binary = {Ops.ADD, Ops.MUL, Ops.IDIV, Ops.MAX, Ops.MOD, Ops.CMPLT, Ops.CMPNE, Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, - Ops.SUB, Ops.FDIV, Ops.POW} + Binary = {Ops.ADD, Ops.MUL, Ops.IDIV, Ops.MAX, Ops.MOD, Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ, + Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW} Ternary = {Ops.WHERE, Ops.MULACC} ALU = set.union(Unary, Binary, Ternary) @@ -88,7 +90,7 @@ class GroupOp: Block = {Ops.BLOCK, Ops.BLOCKEND, Ops.BLOCKSTART} # BinaryOps that can be flipped - Commutative = {Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPNE, Ops.XOR, Ops.AND, Ops.OR} + Commutative = {Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPNE, Ops.CMPEQ, Ops.XOR, Ops.AND, Ops.OR} # BinaryOps where f(f(a,b),c) = f(a,f(b,c)) Associative = {Ops.ADD, Ops.MUL, Ops.AND, Ops.OR, Ops.MAX} diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 7252a76c3c..47b6111977 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -130,7 +130,7 @@ index_pat = UPat(Ops.INDEX, name="idx").or_casted() spec = PatternMatcher([ (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and not x.dtype.local), (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.local), - (UPat(Ops.DEFINE_ACC, src=(UPat.var("c"),), name="x", allow_any_len=True), + (UPat(Ops.DEFINE_REG, src=(UPat.var("c"),), name="x", allow_any_len=True), lambda x,c: all(y.op is Ops.RANGE for y in x.src[1:]) and c.dtype == x.dtype), (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)), @@ -177,7 +177,7 @@ spec = PatternMatcher([ (UPat((Ops.IDIV, Ops.MOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False), (UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype.base == y.dtype.base for y in x.src)), - (UPat(Ops.ASSIGN, src=(UPat((Ops.DEFINE_ACC, Ops.DEFINE_GLOBAL)), UPat())), lambda: True), + (UPat(Ops.ASSIGN, src=(UPat((Ops.DEFINE_REG, Ops.DEFINE_GLOBAL)), UPat())), lambda: True), (UPat(Ops.ENDRANGE, dtype=dtypes.void, src=(UPat(Ops.RANGE),)), lambda: True), # WMMA has a diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 91e0c56e80..6ee9ed5813 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -432,7 +432,7 @@ sym = symbolic_flat+PatternMatcher([ ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32), lambda x,y: y.where(x.cast(dtypes.uint32), UOp.const(dtypes.uint32, 0))), # ** self folding ** - (UPat(Ops.DEFINE_ACC, src=(UPat.var("x"),)), lambda x: x), # a DEFINE_ACC without ranges is a CONST + (UPat(Ops.DEFINE_REG, src=(UPat.var("x"),)), lambda x: x), # a DEFINE_ACC without ranges is a CONST (UPat(Ops.ASSIGN, src=(UPat.cvar(),UPat.var("x"))), lambda x: x), # an ASSIGN to a const is a NOOP # x!=0 -> (bool)x (UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))), diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 7b58df7b83..bb6c9790bb 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -10,7 +10,7 @@ from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileRangeEvent, from tinygrad.dtype import dtypes uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B", - Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_ACC: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B", + Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_REG: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B", Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.VIEW: "#C8F9D4", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",