from typing import Iterator import functools, itertools from dataclasses import dataclass, field, replace from tinygrad.dtype import dtypes, AddrSpace from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes from tinygrad.uop.ops import gate_kernel_sink from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC @dataclass class IndexingContext: realize_map: dict[UOp, None|list[int]] = field(default_factory=dict) non_removable: dict[UOp, None] = field(default_factory=dict) range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict) # loads reachable from each UOp memoized across matches buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict) # create ranges range_idx: Iterator[int] = field(default_factory=itertools.count) def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp: if isinstance(s, UOp) and s.op is Ops.RANGE: return s # if a range has a 1 src, it's the same as UOp.const(0) return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0) ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE, Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM, Ops.LOAD, Ops.CALL, Ops.FUNCTION} def realize(ctx:IndexingContext, tr:UOp) -> None: ctx.realize_map[tr] = None def realize_srcs(ctx:IndexingContext, rb:UOp) -> None: for s in rb.src: if s.base.op not in ALWAYS_CONTIGUOUS: ctx.realize_map[s] = None def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp): # don't realize SLICE when it's the direct source of STORE+AFTER — the target buffer is the output if src.op is Ops.SLICE and src in ctx.realize_map \ and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD): del ctx.realize_map[src] # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce if dest.base in src.backward_slice_with_self: ctx.realize_map[src] = None def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None: for s in c.src[1:]: while s.op is Ops.RESHAPE: s = s.src[0] if s.op not in ALWAYS_CONTIGUOUS: ctx.realize_map[s] = None ctx.non_removable[s] = None pm_generate_realize_map = PatternMatcher([ # realize the inputs of custom kernel calls (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM)),), name="c", allow_any_len=True), realize_custom_kernel_srcs), # always realize (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize), # realize srcs of these (UPat((Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs), # sometimes we need to realize the src of STORE if there's a self-access (UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))), realize_store_after_src), ]) @dataclass(frozen=True) class BufferizeOpts: # on AddrSpace.LOCAL, device is the id device: str|tuple[str, ...]|int|None addrspace: AddrSpace = AddrSpace.GLOBAL removable: bool = True def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]: if x.op not in GroupOp.Broadcastable: return rngs baxes, nleft = broadcast_axes(src.shape, x.shape), len(x.shape)-len(src.shape) return tuple(r.const_like(0) if j in baxes else r for j,r in enumerate(rngs) if j >= nleft) # TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]: if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL, Ops.BIND}: return () if op in GroupOp.Movement|{Ops.INDEX, Ops.SLICE, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1] return src def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]: new_srcs = [] # shape/bound/index args that are not data src should not be indexed data_src_count = len(data_srcs(x.op, x.src)) for i, s in enumerate(x.src): new_src = s src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else () if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}: if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs) elif s in ctx.realize_map: realized_ranges = ctx.realize_map[s] assert isinstance(realized_ranges, list), "realize map must contain range list" closed_ranges = tuple([r for i,r in enumerate(ctx.range_map[s][1]) if i in realized_ranges]) if s.op is Ops.STORE: # add the ends if this is a store new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE]) del ctx.realize_map[s] else: removable = s.op not in ALWAYS_CONTIGUOUS and s not in ctx.non_removable # LOCAL: None in the device assigns it a number later opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \ BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable) new_src = UOp(Ops.STAGE, src=(new_src,)+closed_ranges, arg=opts) if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(src_rngs) if i in realized_ranges]) new_srcs.append(new_src) return new_srcs def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp): if x.op in {Ops.STAGE, Ops.INDEX}: return None return x.replace(src=tuple(create_bufferize_and_index_srcs(ctx, x))) def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp): if x not in ctx.range_map: return None bx = create_bufferize_and_index_based_on_ranges(ctx, x) valid: UOp = UOp.const(True).uprod([r.get_valid() for r in ctx.range_map[x][0]]) return valid.where(bx.src[0], UOp.const(x.dtype.const(0))) def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp): if x.arg[1] == 0: return None bx = create_bufferize_and_index_based_on_ranges(ctx, x) # input ranges new_ranges = list(ctx.range_map[x][0][:x.arg[1]]) return UOp(Ops.REDUCE, src=(bx.src[0],)+tuple(new_ranges), arg=(x.arg[0], 0)) def convert_stack_to_where(ctx:IndexingContext, x:UOp): # only data STACKs: shape tuple STACKs aren't in range_map, the empty shape tuple is void if x not in ctx.range_map or x.dtype == dtypes.void: return None # use the src list directly, a transient STACK of mid-rangeify srcs violates the spec shape rule srcs = create_bufferize_and_index_srcs(ctx, x) r0 = ctx.range_map[x][1][0] ret = srcs[-1] for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret) return ret def remove_movement_op_after_rangeify(ctx:IndexingContext, x:UOp): if x in ctx.range_map or x.src[0].op is Ops.INDEX: return x.src[0] pm_apply_rangeify = PatternMatcher([ # REDUCE(op, axis) -> REDUCE(op) with ranges (UPat(Ops.REDUCE, name="x"), convert_reduce_to_reduce_with_ranges), # PAD -> WHERE (UPat(Ops.PAD, name="x"), convert_pad_to_where_to_keep_behavior_local), # STACK -> WHERE select on the leading range (UPat(Ops.STACK, name="x"), convert_stack_to_where), # finally, apply_rangeify (UPat(GroupOp.All, name="x"), create_bufferize_and_index_based_on_ranges), # remove movement op (UPat(GroupOp.Movement, name="x"), remove_movement_op_after_rangeify), ]) pm_fix_deviceless = PatternMatcher([ (UPat(Ops.STAGE, name="b"), lambda ctx,b: b.replace(arg=replace(b.arg, device=ctx)) if b.arg.addrspace is AddrSpace.GLOBAL and b.arg.device is None else None), ]) @functools.cache def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:UOp) -> UOp: acc:sint = 1 axes_in:list[UOp] = [] for s,src in list(zip(out_shape, urngs.src))[::-1]: axes_in.append(acc*src) acc *= s combined_axes = UOp.const(0).usum(axes_in) axes_out:list[UOp] = [] for s in in_shape[::-1]: axes_out.append(combined_axes % s) combined_axes //= s # this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code return graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_simplify_valid+pm_drop_and_clauses, name="reshape") # this is the definition of the movement ops @functools.cache def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]: match op: case Ops.SHRINK: rngs = tuple(a if off == 0 else a+off for a,(off,_) in zip(rngs, arg)) case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg)) case Ops.FLIP: rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg)) case Ops.EXPAND: rngs = rngs[len(arg):] case Ops.PAD: # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local` # wraps the pad with only the newly added valid rngs = tuple(r if (sz == sh and off == 0) else (r-off).valid(graph_rewrite((r >= off) & (r < (sh+off)), symbolic+pm_simplify_valid, name="pad")) for r,sh,(off,sz) in zip(rngs, in_shape, arg)) case Ops.RESHAPE: sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early sub_array = {r:r.replace(src=r.src[:1], arg=(i, AxisType.PLACEHOLDER)) for i,r in enumerate(sink.ranges)} rngs = _apply_reshape(in_shape, arg, sink.substitute(sub_array)).substitute({v:k for k,v in sub_array.items()}).src case _: raise RuntimeError(f"{op} is not a MovementOp") return rngs @rewrite_group(new_ctx=False) def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: if debug: print("**************************") rctx = IndexingContext() # get ops to realize graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx, name="get realize") # get the consumer map with cpu_profile("consumer map in rangeify", "TINY"): tsink_toposort = tsink.toposort(gate_kernel_sink) consumer_map: dict[UOp, dict[UOp, None]] = {x:{} for x in tsink_toposort} for c in tsink_toposort: for x in data_srcs(c.op, c.src): if x in consumer_map: consumer_map[x][c] = None # explicit rangeify ending_ranges: dict[UOp, list[UOp]] = {} for x in reversed(tsink_toposort): # no ranges on kernels, they are internal if x.op in {Ops.CALL, Ops.FUNCTION, Ops.LINEAR}: continue # AFTER doesn't have range if x.op is Ops.AFTER: continue # treat MSTACK/MSELECT like SINK if x.op in {Ops.MSTACK, Ops.MSELECT}: continue ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], []) # ranges the consumers iterate that this node broadcasts over ended = [rctx.range_map[c][0][i] for c in consumer_map[x] if c in rctx.range_map and c.op in GroupOp.Broadcastable for i in broadcast_axes(x.shape, c.shape)] broadcast_ending_ranges = list(UOp.sink(*ended).ranges) # fusion decision: REDUCE before the broadcast if x.op is Ops.REDUCE: ending_ranges[x] += broadcast_ending_ranges # *** the ranges on the output are # 1. new if this op is realized # 2. from the single consumer if this op only has one consumer # 3. potentially new if this op has 2+ consumers consumer_rngs = [broadcast_rngs(c, x, rctx.range_map[c][0]) for c in consumer_map[x] if c in rctx.range_map] if x in rctx.realize_map: # if this is in the realize_map, we create new ranges (at the output) out_rngs = tuple(rctx.new_range(s) for s in x.shape) # all ranges are ended now ending_ranges[x] = [] # mark all ranges as ended assert rctx.realize_map[x] is None rctx.realize_map[x] = list(range(len(x.shape))) elif len(consumer_rngs) == 0: # if no consumers have ranges and this isn't realized, this doesn't have ranges either. continue elif len(consumer_rngs) == 1: # if this has one consumer, it inherits the ranges from it out_rngs = consumer_rngs[0] elif len(consumer_rngs) > 1: # if this has two consumers, we have to merge the ranges and might create new ones all_rngs: list[tuple[UOp, ...]] = list(zip(*consumer_rngs)) rngs_valids = [] for valid_rngs in all_rngs: local_rngs, valids = zip(*[(r.get_idx(), r.get_valid()) for r in valid_rngs]) rngs_valids.append((local_rngs, valids)) # TODO: in RANGEIFY > 1 all_all_same isn't required all_all_same = all(all_same(local_rngs) for local_rngs,_ in rngs_valids) _out_rngs = [] _realize_axis = [] for i,(local_rngs,valids) in enumerate(rngs_valids): # we compare the ranges without their valids if all_all_same or (PCONTIG and all_same(local_rngs)): # the new valid is the OR of all the children valids minimum_valid = UOp.const(False).usum(valids) _out_rngs.append(graph_rewrite(local_rngs[0].valid(minimum_valid), symbolic, name="minimum_valid")) else: _out_rngs.append(rctx.new_range(x.shape[i])) _realize_axis.append(i) out_rngs = tuple(_out_rngs) # we have to (partially) realize here if there's new ranges if len(_realize_axis): rctx.realize_map[x] = _realize_axis # if this element is a reduce and there's ended ranges, we might have to end some other ranges if len(ending_ranges[x]) and x.op in GroupOp.Elementwise.union({Ops.REDUCE}): _realize_axis = rctx.realize_map.get(x) or [] for i,r in enumerate(out_rngs): if i in _realize_axis: continue if not (PCONTIG > 1) or any(any(rr.arg > e.arg for e in ending_ranges[x]) for rr in r.ranges): _realize_axis.append(i) ending_ranges[x] = [] if len(_realize_axis): rctx.realize_map[x] = _realize_axis out_rngs = tuple([(rctx.new_range(x.shape[i]) if i in _realize_axis else r) for i,r in enumerate(out_rngs)]) ending_ranges[x] += broadcast_ending_ranges # TODO: some ops don't have shape, enable this after the `.st` property is removed #assert len(out_rngs) == len(x.shape), \ # f"shape len mismatch {len(out_rngs)} != {len(x.shape)} on {x.op} with {len(consumer_map[x])} consumers and realize {x in realize_map}" # *** the ranges on the inputs are # 1. swizzled for MovementOps # 2. newly created for REDUCE (tensor graph form with axis) # 3. passed through for everything else rngs = out_rngs # rngs is the input ranges # pylint: disable=possibly-used-before-assignment # apply movement ops if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.marg, rngs) # STACK: the leading range selects the src, srcs get the trailing ranges if x.op is Ops.STACK: rngs = out_rngs[1:] # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do. # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): ending_ranges[x] += list(UOp.sink(*out_rngs[:len(x.marg)]).ranges.keys()) # REDUCE creates ranges for the axes it is reducing if x.op is Ops.REDUCE and x.arg[1]: rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) for s in x.src[0].shape[:x.arg[1]]) + out_rngs if debug: realized_ranges = rctx.realize_map.get(x, None) if x.op is Ops.RESHAPE or len(rngs) != len(out_rngs): disp = render_ranges(rngs, realized=realized_ranges) + " -> " + render_ranges(out_rngs, realized=realized_ranges) else: disp = render_ranges(rngs, out_rngs, realized=realized_ranges) print("***" if x in rctx.realize_map else " ", f"{len(consumer_map[x]):2d} {str(x.op):20s} {str(x._shape):35s} {len(ending_ranges[x]):2d}", disp) # assign to the range map. rngs are the input ranges, out_rngs are the output ranges, from the x op. rctx.range_map[x] = (rngs, out_rngs) # NOTE: SPEC=3 is broken here with shape with Context(SPEC=min(SPEC.value, 2)): tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify") # if a deviceless value must materialize, place it on the sink device tsink = graph_rewrite(tsink, pm_fix_deviceless, ctx=tsink.device, name="add device to deviceless") return tsink, rctx def render_ranges(*rngs_list, realized) -> str: disp = [] for i, rs in enumerate(zip(*[[r.render() for r in rngs] for rngs in rngs_list])): rng = rs[0] if all_same(rs) else " -> ".join(rs) if realized is not None and i in realized: rng = colored(rng, "yellow") disp.append("["+rng+"]") return ''.join(disp)