From cbcc1c20eb09a1342f6581cfbb99632bade982a8 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 2 Dec 2024 13:43:09 +0800 Subject: [PATCH] second try at block linearize (#7892) * second try at block linearize * weeee, works for lil matmul * it's so beautiful * test tiny passes * fix bugs * combine matching BLOCKENDS * wrapping * test lin failures passes * those failures were fake * flip sort order * fix ptx tests * deal with store better * dumb ptx fix * expect less * reduce lines * reduce lines * less lines and cleaner * no defaultdict * tighter * simpler block_parent_count --- test/external/speed_v_theoretical.py | 2 +- tinygrad/codegen/linearize.py | 228 ++++++++++++++++++--------- 2 files changed, 151 insertions(+), 79 deletions(-) diff --git a/test/external/speed_v_theoretical.py b/test/external/speed_v_theoretical.py index 0dc2284ed9..88bfc7bfa5 100644 --- a/test/external/speed_v_theoretical.py +++ b/test/external/speed_v_theoretical.py @@ -88,7 +88,7 @@ class TestKernelSpeed(unittest.TestCase): # def test_gemm_1024(self): self._test_matmul(1024, nv_tflops=8, amd_tflops=7) # def test_gemm_2048(self): self._test_matmul(2048, nv_tflops=50, amd_tflops=30) def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=95, amd_tflops=70) - def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=130, amd_tflops=70) + def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=125, amd_tflops=70) def test_gemv_16384_4096(self): self._test_matmul(16384, 4096, 1, nv_gbs=430, amd_gbs=400) def test_gemv_4096_16384(self): self._test_matmul(4096, 16384, 1, nv_gbs=430, amd_gbs=380) # AMD was flaky at 400 diff --git a/tinygrad/codegen/linearize.py b/tinygrad/codegen/linearize.py index 356835b7bc..aab5ad3e2b 100644 --- a/tinygrad/codegen/linearize.py +++ b/tinygrad/codegen/linearize.py @@ -1,92 +1,164 @@ -from typing import List, Set, Dict, Tuple -import functools, heapq -from tinygrad.ops import type_verify, END_FOR_UOP, UOp, Ops, GroupOp -from tinygrad.dtype import dtypes -from tinygrad.helpers import DEBUG +from typing import List, Dict, Tuple +import functools, collections +from tinygrad.ops import type_verify, UOp, Ops, PatternMatcher, UPat, graph_rewrite +from tinygrad.dtype import dtypes, PtrDType +from tinygrad.helpers import dedup, flatten, partition -def get_children_dfs(u:UOp, children:Dict[UOp, List[UOp]], srcs:Dict[UOp, Dict[UOp, None]], in_degree:Dict[UOp, int]): - if u in children: return srcs[u] - srcs[u] = {} - children[u] = [] - for x in u.src: - srcs[u].update(get_children_dfs(x, children, srcs, in_degree)) - if x.op is Ops.RANGE and x.arg[1]: srcs[u][x] = None - children[x].append(u) - in_degree[u] = len(u.src) - return srcs[u] +DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST, + Ops.BLOCK, Ops.BLOCKEND, Ops.BLOCKFORK, Ops.BLOCKSTART} + +def disp(y:UOp) -> str: + if y.op is Ops.BLOCKSTART: return "w"+disp(y.src[0]) + if y.op is Ops.IF: return f'IF{id(y)}' + if y.op is Ops.RANGE: return str(y.arg[0]) + return "" + +class BasicBlock: + def __init__(self, ctx, lst, end=None): + self.ctx, self.lst, self.end = ctx, lst, end + def __repr__(self): + return f"{(str(disp(self.end))+' ') if self.end is not None else ''}"+\ + f"{[disp(y) for y in self.ctx]} {len(self.lst)}" + "\n" + '\n'.join([str(x.op) for x in self.lst]) + +def append_to_block(ctx, x:UOp): + block_ctxs, children = ctx + new_srcs: List[UOp] = [] + to_append: List[UOp] = [] + new_blocks: Dict[Tuple[UOp, ...], List[UOp]] = {} + in_this_block = set(x.arg.lst) + for u in x.src: + if u.op in DONT_PLACE_IN_BLOCK or len([y for y in children[u] if y not in in_this_block]) > 0: + # if it's a fork or not placed, we don't place it + new_srcs.append(u) + elif (block_ctx:=block_ctxs[u]) == x.arg.ctx: + # if it's the same context, we place the UOp in this block and append the parents to it's srcs + new_srcs += list(u.src) + to_append.append(u) + else: + # otherwise, we create a new block with this UOp + new_blocks.setdefault(block_ctx, []).append(u) + if len(to_append) == 0 and len(new_blocks) == 0: return None + + for rng,lst in new_blocks.items(): + new_block = UOp(Ops.BLOCK, dtypes.void, tuple(dedup(flatten(y.src for y in lst))), BasicBlock(rng, lst)) + lrng = list(rng) + for r in rng[::-1]: + if r not in x.arg.ctx and r.op is not Ops.BLOCKSTART: + lrng.remove(r) + new_block = UOp(Ops.BLOCKEND, src=(new_block,), arg=BasicBlock(lrng[:], [UOp(Ops.ENDIF if r.op is Ops.IF else Ops.ENDRANGE, src=(r,))], r)) + new_srcs.append(new_block) + return UOp(Ops.BLOCK, dtypes.void, tuple(dedup(new_srcs)), BasicBlock(x.arg.ctx, to_append+x.arg.lst)) + +make_basic_blocks = PatternMatcher([ + (UPat(Ops.SINK, name="x"), lambda x: UOp(Ops.BLOCK, src=x.src, arg=BasicBlock([], [x]))), + (UPat(Ops.BLOCK, name="x"), append_to_block), +]) + +def block_merge(ctx, x:UOp): + # ctx is children here + if x.op is Ops.BLOCKEND: + # if it's a BLOCKEND, see if we are done with placement. if all the children of the range are in here + in_this_block = set(x.arg.lst) + if len([y for y in ctx[x.arg.end] if y not in in_this_block]) == 0: + # find the parent block that has the BLOCKSTART in the ctx + parent_blocks = [y for y in x.src if y.op is Ops.BLOCK and UOp(Ops.BLOCKSTART, src=(x.arg.end,)) in y.arg.ctx] + if len(parent_blocks) == 1: + parent_block = parent_blocks[0] + # 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) + return UOp(Ops.BLOCK, dtypes.void, tuple(y for y in x.src if y is not parent_block)+parent_block.src, + BasicBlock([y for y in x.arg.ctx if y is not x.arg.end], early_ops+parent_block.arg.lst+late_ops)) + assert not len(parent_blocks) + + new_srcs: List[UOp] = [] + to_append: List[UOp] = [] + new_ctx = list(x.arg.ctx[:]) + placed = set() + for u in x.src: + if u.op is Ops.BLOCK and (tuple(u.arg.ctx) == tuple(x.arg.ctx) or (x.arg.end is not None and x.arg.end in u.arg.ctx)): + # NOTE: this can't appear in srcs twice or it would be a BLOCKFORK + new_ctx += u.arg.ctx + new_srcs += list(u.src) + to_append += u.arg.lst + elif u.op is Ops.BLOCKFORK and len([y for y in x.src if y is u]) == u.arg: # block fork appears # of times in srcs + if u not in placed: + new_srcs += list(u.src) + placed.add(u) + else: + # keep it in srcs + new_srcs.append(u) + if len(to_append) == 0 and len(placed) == 0: return None + return UOp(x.op, dtypes.void, tuple(new_srcs), BasicBlock(dedup(new_ctx), to_append+x.arg.lst, x.arg.end)) + +pm_block_merge = PatternMatcher([(UPat((Ops.BLOCKEND, Ops.BLOCK), name="x"), block_merge),]) def linearize_uop(sink:UOp, skip_check:bool=not __debug__) -> List[UOp]: assert sink.op is Ops.SINK, f"sink isn't sink, it's {sink.op}" - # filter nodes that don't link to a sink - # BFS toposort + + @functools.lru_cache(None) + def get_block_ctx(x:UOp) -> Tuple[UOp, ...]: + ret: List[UOp] = [] + for u in x.src: + if u.op in {Ops.RANGE, Ops.IF}: ret.append(u) + # don't flow (fully) through assign and store + elif u.op is Ops.STORE: + # ugh, deal with non-reduce locals. probably wrong + if isinstance(u.src[0].dtype, PtrDType) and u.src[0].dtype.local: + idx_context, store_context = get_block_ctx(u.src[0]), get_block_ctx(u) + ret += [x for x in store_context if x not in idx_context and x.op is Ops.RANGE] + elif u.op is Ops.ASSIGN: + # flow though assign, but remove the ranges used in the assign + assert u.src[0].op is Ops.DEFINE_ACC + ret += [x for x in get_block_ctx(u.src[1]) if x not in u.src[0].src[1:]] + else: + # flow though everything else + ret += get_block_ctx(u) + return tuple(dedup(sorted(ret, key=lambda x: x.tuplize))) + + # get children and all block contexts + block_ctxs: Dict[UOp, Tuple[UOp, ...]] = {} children: Dict[UOp, List[UOp]] = {} - range_srcs: Dict[UOp, Dict[UOp, None]] = {} - in_degree: Dict[UOp, int] = {} - get_children_dfs(sink, children, range_srcs, in_degree) + for u in sink.sparents: + for s in u.src: children.setdefault(s, []).append(u) + this_block_ctx = get_block_ctx(u) + block_ctxs[u] = ((UOp(Ops.BLOCKSTART, src=(u,)),) + this_block_ctx) if u.op in {Ops.IF, Ops.RANGE} else this_block_ctx - @functools.lru_cache(None) - def get_recursive_children(x:UOp, end:Ops, include_self=False) -> Set[UOp]: - if x.op is Ops.SINK: return set() - return set.union({x} if include_self else set(), *([get_recursive_children(u, end, True) for u in children[x] if x.op is not end])) + # TODO: there's probably a clever way to remove this while loop + while 1: + sink = graph_rewrite(sink, make_basic_blocks, ctx=(block_ctxs, children)) - # scope children impact the toposort and END* insertion - scope_children = {p:get_recursive_children(p, END_FOR_UOP[p.op][0]) for p in reversed(in_degree) if p.op in END_FOR_UOP} - range_phi = {r:[p for p in scope_children[r] if p.op is Ops.ASSIGN] for r in scope_children if r.op is Ops.RANGE} + # add BLOCKFORK (slow!) + block_parent_count = collections.Counter(flatten([x.src for x in sink.sparents if x.op is Ops.BLOCK])) + non_block_parents = flatten([x.src for x in sink.sparents if x.op is not Ops.BLOCK]) + forks = {} + for u,child_count in block_parent_count.items(): + if u.op not in DONT_PLACE_IN_BLOCK and child_count > 1 and u not in non_block_parents: + forks[u] = UOp(Ops.BLOCKFORK, src=(UOp(Ops.BLOCK, src=u.src, arg=BasicBlock(block_ctxs[u], [u])),), arg=child_count) - # assign priorities - def get_priority(u:UOp): - priority = 0 - # prefer ranges that depend on the least number of independent ranges - if u.op is Ops.RANGE and u.arg[1]: - priority += u.arg[0] - for p in range_phi[u]: - priority += 10000*len([r for r in range_srcs[p] if not any(i in range_phi[u] for i in range_phi[r])]) - elif u.op is Ops.CONST: - # place consts first here, they don't do anything and it can cause issues with DEFINE_ACC - priority -= 100000000000 - else: - # prefer uops that are loop children - priority -= sum([(l.arg[0]+1) + 1000*l.arg[1] for l,ss in scope_children.items() if l.op is Ops.RANGE and u in ss]) - if u.op is Ops.IF and len(u.src) == 1: priority += 10000000 # if penalty - return priority - priorities:Dict[UOp, int] = {u:get_priority(u) for u in children} + if not len(forks): break + sink = sink.substitute(forks) - # prevent priority inversion - @functools.lru_cache(None) - def fix_priority(u:UOp, lowest_priority): - if u.op in {Ops.CAST, Ops.BITCAST, *GroupOp.ALU, Ops.VECTORIZE, Ops.GEP, Ops.SPECIAL, Ops.DEFINE_LOCAL, Ops.LOAD}: - priorities[u] = min(priorities[u], lowest_priority) - if u.op is Ops.LOAD: priorities[u] += 100 # load penalty (here) - for x in u.src: fix_priority(x, priorities[u]) - fix_priority(sink, 0) + # combine matching BLOCKENDS + blockends_to_arg: Dict[UOp, List[UOp]] = {} + for be in sink.sparents: + if be.op is Ops.BLOCKEND: blockends_to_arg.setdefault(be.arg.end, []).append(be) + new_forks = {} + for k,v in blockends_to_arg.items(): + # NOTE: if any BLOCKEND is the parent of any other with the same arg, this algo fails + if len(v) > 1: + new_blockend = UOp(Ops.BLOCKEND, src=tuple(flatten(x.src for x in v)), arg=BasicBlock(dedup(flatten([y.arg.ctx for y in v])), v[0].arg.lst, k)) + out = UOp(Ops.BLOCKFORK, src=(new_blockend,), arg=len(v)) + for u in v: new_forks[u] = out + sink = sink.substitute(new_forks) - # NOTE: the compare should never make it all the way to u - queue:List[Tuple[int, Tuple, UOp]] = [] - def push(u:UOp): heapq.heappush(queue, (priorities[u], u.tuplize, u)) + # final rewrite to merge all blocks into one + sink = graph_rewrite(sink, pm_block_merge, ctx=children) - for u in children: - if in_degree[u] == 0: push(u) - - scope_end: Dict[UOp, UOp] = {} - _uops: List[UOp] = [] - while queue: - p,_,x = heapq.heappop(queue) - if DEBUG >= 7: print(f"{p:5d}", x.op, x.dtype, x.arg) - if x in scope_children: scope_end[x] = x - if x.op is Ops.DEFINE_ACC: - idx = min([_uops.index(l) for l in x.src if l.op is Ops.RANGE]) - _uops.insert(idx, x) - else: _uops.append(x) - for u, ss in scope_children.items(): - if x in ss: - ss.remove(x) - if len(ss) == 0: scope_end[u] = x - for u in children[x]: - in_degree[u] -= 1 - if in_degree[u] == 0: push(u) - - # end scopes in toposort order - for u, x in scope_end.items(): _uops.insert(_uops.index(x)+1, UOp(END_FOR_UOP[u.op][1], dtypes.void, (u,))) + # there should just be one block left, with a few parents with 0 srcs + assert sink.op is Ops.BLOCK + _uops = sorted(dedup(sink.src), key=lambda x: x.tuplize) + assert all(len(x.src) == 0 and x.op not in {Ops.BLOCK, Ops.BLOCKSTART, Ops.BLOCKEND, Ops.BLOCKFORK} for x in _uops) + _uops += sink.arg.lst # sanity checks (NOTE: these can cause things to be skipped in BEAM) if not skip_check: type_verify(_uops)