geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

Replace ASSIGN with STORE in UPat compiler (#11065)

Browse Source
This commit is contained in:
quortus
2025-07-02 19:15:43 -07:00
committed by GitHub
parent d049639221
commit a937ac80dc
1 changed files with 21 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/uop/upat.py
+21 -21
View File
@@ -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):