diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 258c51a317..c8e74fd2d7 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -1,6 +1,6 @@ import unittest, itertools -from tinygrad.codegen.late.devectorizer import indexing_simplify +from tinygrad.codegen.late.coalese import indexing_simplify from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops, graph_rewrite from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load @@ -495,7 +495,7 @@ class TestImageSimplification(unittest.TestCase): class TestDropTrueGate(unittest.TestCase): def test_drop_true_gate_on_index(self): # test that INDEX with a constant True valid gets simplified to drop the valid - from tinygrad.codegen.late.devectorizer import indexing_simplify + from tinygrad.codegen.late.coalese import indexing_simplify from tinygrad.uop.ops import graph_rewrite from tinygrad.uop.symbolic import sym buf = UOp.param(0, dtypes.int.ptr()) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 38e473df1e..520aae7b28 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -16,7 +16,7 @@ from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_ from tinygrad.codegen.decomp.dtype import pm_dtype_decomps from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns -from tinygrad.codegen.late.devectorizer import indexing_simplify +from tinygrad.codegen.late.coalese import indexing_simplify from tinygrad.codegen.opt.postrange import apply_opts from tinygrad.codegen.late.gater import pm_move_gates_from_index from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse diff --git a/tinygrad/codegen/late/coalese.py b/tinygrad/codegen/late/coalese.py index 3ba440f7ce..de22a0fb15 100644 --- a/tinygrad/codegen/late/coalese.py +++ b/tinygrad/codegen/late/coalese.py @@ -1,11 +1,70 @@ from typing import Any -import itertools +import itertools, functools from collections import defaultdict -from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType -from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat -from tinygrad.helpers import getenv, IMAGE +from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType, DType +from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp +from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate +from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv from tinygrad.renderer import Renderer -from tinygrad.codegen.late.devectorizer import image_valid_dims, _drop_valid_stmts, uop_given_valid + +# ***** image load valid simplification ***** + +@functools.cache +def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]: + # can drop valid if idx is out of bound when valid is False + drop_stmt = [] + for i,stmt in enumerate(valid.split_uop(Ops.AND)): + if (res:=parse_valid(stmt)) is None: continue + X, is_upper_bound, c = res + + # 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 X.split_uop(Ops.ADD)): + testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx) + if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0: + drop_stmt.append(stmt) + continue + + # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1] + lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1) + if lo <= hi: + fake = UOp.variable(f"fake{i}", lo, hi, X.dtype) + for coord,b in zip(idx.src, (width, height)): + rw = coord.substitute({X:fake}).simplify() + if rw.vmin >= b or rw.vmax < 0: + drop_stmt.append(stmt) + break + return drop_stmt + +def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: + idx = uop_given_valid(valid, start_idx) + return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True) + +def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None: + if not isinstance(buf.dtype, ImageDType): return None + start_idx = idx_x._stack(idx_y) + idx = uop_given_valid(valid, start_idx) + drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1]) + + if not drop_stmt and idx is start_idx: return None + new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None + idx_y, idx_x = idx.index(1), idx.index(0) + return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True) + +indexing_simplify = PatternMatcher([ + # image load valid idx simplification + (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)), + (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)), + UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load), +]) + +# get list of (height, width) that do not require pitch padding +def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]: + if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return [] + MAXW, pxls = 16384, size // 4 + if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return [] + # height=1 images just need to abide by alignment requirements in bytes, not pixels! + if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)] + return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0] def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None: shapes, ren = ctx diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py deleted file mode 100644 index 83e2c74b3b..0000000000 --- a/tinygrad/codegen/late/devectorizer.py +++ /dev/null @@ -1,66 +0,0 @@ -import functools -from tinygrad.dtype import dtypes, ImageDType, DType, Invalid -from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp -from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate -from tinygrad.helpers import OSX, ceildiv - -# ***** image load valid simplification ***** - -@functools.cache -def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]: - # can drop valid if idx is out of bound when valid is False - drop_stmt = [] - for i,stmt in enumerate(valid.split_uop(Ops.AND)): - if (res:=parse_valid(stmt)) is None: continue - X, is_upper_bound, c = res - - # 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 X.split_uop(Ops.ADD)): - testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx) - if testidx.index(0).vmax < 0 or testidx.index(1).vmax < 0: - drop_stmt.append(stmt) - continue - - # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1] - lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1) - if lo <= hi: - fake = UOp.variable(f"fake{i}", lo, hi, X.dtype) - for coord,b in zip(idx.src, (width, height)): - rw = coord.substitute({X:fake}).simplify() - if rw.vmin >= b or rw.vmax < 0: - drop_stmt.append(stmt) - break - return drop_stmt - -def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None: - idx = uop_given_valid(valid, start_idx) - return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True) - -def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None: - if not isinstance(buf.dtype, ImageDType): return None - start_idx = idx_x._stack(idx_y) - idx = uop_given_valid(valid, start_idx) - drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1]) - - if not drop_stmt and idx is start_idx: return None - new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None - idx_y, idx_x = idx.index(1), idx.index(0) - return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True) - -indexing_simplify = PatternMatcher([ - # image load valid idx simplification - (UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)), - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)), - UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load), -]) - -# ***** load/store grouping ***** - -# get list of (height, width) that do not require pitch padding -def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]: - if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return [] - MAXW, pxls = 16384, size // 4 - if base not in (dtypes.half, dtypes.float) or size > 4*MAXW*MAXW: return [] - # height=1 images just need to abide by alignment requirements in bytes, not pixels! - if size % (ALIGN * 4) != 0: return [] if (base.itemsize * size) % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)] - return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0] diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 812b91913b..36cc990679 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -3,7 +3,7 @@ from tinygrad.codegen.opt import Opt, OptOps, KernelOptError from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, IMAGE from tinygrad.dtype import PtrDType from tinygrad.uop.ops import Ops, resolve, AxisType -from tinygrad.codegen.late.devectorizer import image_valid_dims +from tinygrad.codegen.late.coalese import image_valid_dims from tinygrad.codegen.opt.postrange import Scheduler def hand_coded_optimizations(k:Scheduler) -> Scheduler: