From a937ac80dc2e06cfdb92ef088e2261c8618ac971 Mon Sep 17 00:00:00 2001 From: quortus <156855065+quortus@users.noreply.github.com> Date: Thu, 3 Jul 2025 04:15:43 +0200 Subject: [PATCH] Replace ASSIGN with STORE in UPat compiler (#11065) --- tinygrad/uop/upat.py | 42 +++++++++++++++++++++--------------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/tinygrad/uop/upat.py b/tinygrad/uop/upat.py index 090de21357..5c5f221d31 100644 --- a/tinygrad/uop/upat.py +++ b/tinygrad/uop/upat.py @@ -21,7 +21,7 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp: else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.arg)), arg="{0}.arg == {1}")) if self.strict_length or self.required_len > 0: and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len)))) - if self.name is not None: and_clause.append(UOp(Ops.ASSIGN, src=(UOp(Ops.DEFINE_VAR, arg=self.name), base))) + if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.DEFINE_VAR, arg=self.name), base))) if self.dtype is not None: if len(self.dtype) > 1: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.dtype))), arg="({0}.dtype in {1} or {0}.dtype._scalar in {1})")) @@ -66,26 +66,26 @@ def do_process_and(a:UOp) -> UOp|None: or_clause = [UOp(Ops.OR, src=tuple([UOp(Ops.AND, src=x) for x in itertools.product(*[x.src for x in or_clause])]))] found = True - # handle assigns - assigns, new_src = partition(new_src, lambda x: x.op is Ops.ASSIGN) - if len(assigns): + # handle stores + stores, new_src = partition(new_src, lambda x: x.op is Ops.STORE) + if len(stores): if len(or_clause): - # push assigns to the top if we have an or_clause + # push stores to the top if we have an or_clause assert len(or_clause) == 1 and all(x.op is Ops.AND for x in or_clause[0].src) - or_clause = [UOp(Ops.OR, src=tuple([x.replace(src=x.src+tuple(assigns)) for x in or_clause[0].src]))] + or_clause = [UOp(Ops.OR, src=tuple([x.replace(src=x.src+tuple(stores)) for x in or_clause[0].src]))] found = True else: - # check for duplicate assigns - dict_assigns: dict[UOp, UOp] = {} - for a in assigns: - if a.src[0] in dict_assigns: - # duplicate assign is a compare - new_src.append(UOp(Ops.CMPNE, src=(dict_assigns[a.src[0]], a.src[1]))) + # check for duplicate stores + dict_stores: dict[UOp, UOp] = {} + for a in stores: + if a.src[0] in dict_stores: + # duplicate store is a compare + new_src.append(UOp(Ops.CMPNE, src=(dict_stores[a.src[0]], a.src[1]))) found = True else: - dict_assigns[a.src[0]] = a.src[1] - # put the assigns back - for k,v in dict_assigns.items(): new_src.append(UOp(Ops.ASSIGN, src=(k,v))) + dict_stores[a.src[0]] = a.src[1] + # put the stores back + for k,v in dict_stores.items(): new_src.append(UOp(Ops.STORE, src=(k,v))) # reassemble, if there's any deduping to do, do it if len(dretand:=dedup(new_src+or_clause)) != len(new_src)+len(or_clause): found = True @@ -115,25 +115,25 @@ pm_renderer = PatternMatcher([ def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]: assert x.op is Ops.AND - and_pieces, assign_pieces = [], [] + and_pieces, store_pieces = [], [] or_pieces: list[str] = [] for s in x.src: if s.op is Ops.OR: assert len(or_pieces) == 0 and len(s.src) >= 1 for ss in s.src: or_pieces.extend(_final_render(ss, has_ctx, depth+1)) - elif s.op is Ops.ASSIGN: + elif s.op is Ops.STORE: assert s.src[0].op is Ops.DEFINE_VAR and s.src[1].op is Ops.NOOP - assign_pieces.append(f"{s.src[0].arg}={s.src[1].arg}") + store_pieces.append(f"{s.src[0].arg}={s.src[1].arg}") elif s.op is Ops.NOOP: and_pieces.append(s.arg) else: raise UPatCompileError(f"can't compile this {s}") # if we have an or, render it if len(or_pieces): - assert len(assign_pieces) == 0 + assert len(store_pieces) == 0 and_clause = ' and '.join(and_pieces) return [f"{' '*depth}if {and_clause if len(and_clause) else 'True'}:"] + or_pieces # if we don't, this is a final return - assign_clause = ', '.join((["ctx=ctx"] if has_ctx else [])+assign_pieces) - and_clause = ' and '.join(and_pieces + [f"(_ret:=_fxn({assign_clause})) is not None"]) + store_clause = ', '.join((["ctx=ctx"] if has_ctx else [])+store_pieces) + and_clause = ' and '.join(and_pieces + [f"(_ret:=_fxn({store_clause})) is not None"]) return [f"{' '*depth}if {and_clause}: return _ret"] def _get_code(self:UPat, has_ctx:bool):