From 28e06d2d4461e9f625cd9c0724afa54a3a29317f Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 7 Apr 2025 11:28:14 +0800 Subject: [PATCH] minor cleanups from patternmatcher [pr] (#9756) --- tinygrad/codegen/devectorizer.py | 106 +++++++++++++++---------------- tinygrad/ops.py | 3 +- 2 files changed, 55 insertions(+), 54 deletions(-) diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index 82881f6ce2..e4fc9407d3 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -9,6 +9,59 @@ from tinygrad.helpers import getenv, flatten, TRANSCENDENTAL, AMX, prod, DEVECTO from tinygrad.codegen.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES from tinygrad.renderer import Renderer +# ***** image load valid simplification ***** + +def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: + if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.const_like(0) + if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx, valid) + + # wait for it to be image indexed before running simplification + if start_idx.dtype.count != 2: return None + + # can drop valid if idx is out of bound when valid is False + drop_stmt = [] + for stmt in split_uop(valid, Ops.AND): + X, is_upper_bound, c = parse_valid(stmt) + + # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i + if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in split_uop(X, Ops.ADD)): + testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), split_uop(X, Ops.ADD), idx) + testidx = testidx.simplify() + if testidx.gep(0).vmax < 0 or testidx.gep(1).vmax < 0: + drop_stmt.append(stmt) + continue + + # if X <= c, check if it's out of bound when X = c+1 + # if X >= c, check if it's out of bound when X = c-1 + test_value = c + 1 if is_upper_bound else c - 1 + for i,b in zip(idx.src, (buf.dtype.shape[1], buf.dtype.shape[0])): + if i.is_increasing(): + rw = i.substitute({X:X.const_like(test_value)}).simplify() + if rw.vmin >= b or rw.vmax < 0: + drop_stmt.append(stmt) + break + + if not drop_stmt and idx is start_idx: return None + new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in split_uop(valid, Ops.AND) if s not in drop_stmt]) else None + return buf.index(idx, new_valid) + +def delete_redundant_gates(buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None: + if store_gate not in [gate.src[0] for gate in val.toposort if gate.op is Ops.IF]: return None + # remove the gate from the index + return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val) + +load_store_indexing = PatternMatcher([ + # simplify valid + (UPat(Ops.AND, name="valid"), simplify_valid), + # image load valid idx simplification + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat.var("valid"))), simplify_valid_load), + # index True is just Index + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat(Ops.CONST, arg=True))), lambda buf,start_idx: buf.index(start_idx)), + # delete_redundant_gates (after expand) + (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")), + UPat.var("val"))), delete_redundant_gates), +]) + # ***** load/store grouping ***** def expand_index(buf:UOp, vec:UOp, mask:UOp|None=None): @@ -84,42 +137,6 @@ load_store_folding = PatternMatcher([ (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data"))), cat_after_store), ]) -# ***** image load valid simplification ***** - -def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: - if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.const_like(0) - if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx, valid) - - # wait for it to be image indexed before running simplification - if start_idx.dtype.count != 2: return None - - # can drop valid if idx is out of bound when valid is False - drop_stmt = [] - for stmt in split_uop(valid, Ops.AND): - X, is_upper_bound, c = parse_valid(stmt) - - # for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i - if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in split_uop(X, Ops.ADD)): - testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), split_uop(X, Ops.ADD), idx) - testidx = testidx.simplify() - if testidx.gep(0).vmax < 0 or testidx.gep(1).vmax < 0: - drop_stmt.append(stmt) - continue - - # if X <= c, check if it's out of bound when X = c+1 - # if X >= c, check if it's out of bound when X = c-1 - test_value = c + 1 if is_upper_bound else c - 1 - for i,b in zip(idx.src, (buf.dtype.shape[1], buf.dtype.shape[0])): - if i.is_increasing(): - rw = i.substitute({X:X.const_like(test_value)}).simplify() - if rw.vmin >= b or rw.vmax < 0: - drop_stmt.append(stmt) - break - - if not drop_stmt and idx is start_idx: return None - new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in split_uop(valid, Ops.AND) if s not in drop_stmt]) else None - return buf.index(idx, new_valid) - # ***** optional patterns ***** powers_of_two = {2**i:i for i in range(64)} @@ -248,23 +265,6 @@ devectorize = PatternMatcher([ (UPat(Ops.DEFINE_ACC, name="acc"), no_vectorized_acc), ]) -def delete_redundant_gates(buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None: - if store_gate not in [gate.src[0] for gate in val.toposort if gate.op is Ops.IF]: return None - # remove the gate from the index - return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val) - -load_store_indexing = PatternMatcher([ - # simplify valid - (UPat(Ops.AND, name="valid"), simplify_valid), - # image load valid idx simplification - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat.var("valid"))), simplify_valid_load), - # index True is just Index - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat(Ops.CONST, arg=True))), lambda buf,start_idx: buf.index(start_idx)), - # delete_redundant_gates (after expand) - (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")), - UPat.var("val"))), delete_redundant_gates), -]) - pm_render = PatternMatcher([ # for rendering, we use explicit VECTORIZE (UPat(Ops.CONST, name='c'), diff --git a/tinygrad/ops.py b/tinygrad/ops.py index b946b4f87a..1cd7340abe 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -733,6 +733,7 @@ class UPat(MathTrait): upat_match = [src] if isinstance(src, UPat) else ([] if src is None else self.src[0]) self.early_reject = {pp.op[0] for pp in upat_match if pp.op is not None and len(pp.op) == 1} + def __reduce__(self): return UPat,(self.op, self.dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject) def named(self, name:str): return UPat(self.op, self.dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject) @staticmethod @@ -771,7 +772,7 @@ class UPat(MathTrait): def rep(x): form = "UPat(%s, %s, name=%s, dtype=%s, allow_any_len=%s, src=%s)" return form % (None if x.op is None else ('(%s)'%', '.join(map(str, x.op))), x.arg, repr(x.name), - set(x.dtype) if x.dtype else None, x.allowed_len == 0, "[%s]" if x.src and len(x.src)>1 else "(%s)") + set(x.dtype) if x.dtype else None, not x.strict_length, "[%s]" if x.src and len(x.src)>1 else ("(%s)" if x.src else "%s")) return pretty_print(self, rep, srcfn=lambda x:None if x.src is None else [next(x.src[0])] if isinstance(x.src[0], itertools.repeat) else x.src[0]) def match(self:UPat, uop:UOp, store:dict[str, UOp]) -> list[dict[str, UOp]]: