forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ 0798119
This commit is contained in:
IQ.Lvbs history cleanup
286
tinygrad_repo/tinygrad/schedule/indexing.py
Normal file
286
tinygrad_repo/tinygrad/schedule/indexing.py
Normal file
@@ -0,0 +1,286 @@
|
||||
from typing import Iterator
|
||||
import functools, itertools
|
||||
from dataclasses import dataclass, field
|
||||
from tinygrad.dtype import dtypes, AddrSpace
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
|
||||
from tinygrad.uop.ops import consumer_map_from_toposort, 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, IMAGE
|
||||
|
||||
ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.COPY, Ops.BUFFER, Ops.SLICE,
|
||||
Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
|
||||
Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.CALL, Ops.FUNCTION}
|
||||
|
||||
def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
|
||||
|
||||
def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
|
||||
for s in rb.src:
|
||||
if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
|
||||
|
||||
def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
|
||||
# don't realize COPY/SLICE when they are the direct source of STORE+AFTER — the target buffer is the output
|
||||
if src.op in {Ops.COPY, Ops.SLICE} and src in ctx \
|
||||
and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
|
||||
del ctx[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[src] = None
|
||||
|
||||
pm_generate_realize_map = PatternMatcher([
|
||||
# always realize
|
||||
(UPat({Ops.COPY, Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
|
||||
# realize srcs of these
|
||||
(UPat((Ops.COPY, 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
|
||||
|
||||
@dataclass
|
||||
class IndexingContext:
|
||||
realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
|
||||
range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[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.LOOP) -> 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(dtypes.weakint, 0)
|
||||
return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.weakint, 0)
|
||||
|
||||
def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
|
||||
if x.op in {Ops.STAGE, Ops.INDEX}: return None
|
||||
new_srcs = []
|
||||
for s in x.src:
|
||||
new_src = s
|
||||
if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
|
||||
if x in ctx.range_map: new_src = new_src.index(*ctx.range_map[x][0])
|
||||
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:
|
||||
# the Bufferize before a COPY is not removable. there should be a better way to do this
|
||||
removable = x.op is not Ops.COPY and s.op not in ALWAYS_CONTIGUOUS
|
||||
# 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, s.dtype, src=(new_src,)+closed_ranges, arg=opts)
|
||||
if x in ctx.range_map: new_src = new_src.index(*[r for i,r in enumerate(ctx.range_map[x][0]) if i in realized_ranges])
|
||||
new_srcs.append(new_src)
|
||||
# NOTE: do we need this?
|
||||
return x.replace(src=tns) if x.src != (tns:=tuple(new_srcs)) else None
|
||||
|
||||
def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
|
||||
if x not in ctx.range_map: return None
|
||||
valid: UOp = UOp.const(dtypes.bool, True).uprod([r.get_valid() for r in ctx.range_map[x][0]])
|
||||
ret = valid.where(x.src[0], UOp.const(x.dtype, 0))
|
||||
ctx.range_map[ret] = ctx.range_map[x]
|
||||
return ret
|
||||
|
||||
def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
|
||||
if len(x.arg[1]) == 0: return None
|
||||
# input ranges
|
||||
new_ranges = [r for i,r in enumerate(ctx.range_map[x][0]) if i in x.arg[1]]
|
||||
ret = UOp(Ops.REDUCE, x.dtype, src=(x.src[0],)+tuple(new_ranges), arg=(x.arg[0], ()))
|
||||
ctx.range_map[ret] = ctx.range_map[x]
|
||||
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),
|
||||
# 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),
|
||||
])
|
||||
|
||||
@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(dtypes.weakint, 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 = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, 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 graph_rewrite((r >= off) & (r < (sh+off)),
|
||||
symbolic+pm_simplify_valid, name="pad").where(r-off, UOp.invalid()) 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:UOp.range(r.src[0], 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
|
||||
|
||||
@profile_matches
|
||||
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.realize_map, name="get realize")
|
||||
|
||||
# get the consumer map
|
||||
with cpu_profile("consumer map in rangeify", "TINY"):
|
||||
consumer_map = consumer_map_from_toposort(tsink_toposort:=tsink.toposort(gate_kernel_sink))
|
||||
|
||||
# explicit rangeify
|
||||
ending_ranges: dict[UOp, list[UOp]] = {}
|
||||
reduce_bodies = {u:set(u.src[0].toposort()) for u in tsink_toposort if u.op is Ops.REDUCE and len(u.arg[1])} if IMAGE else {}
|
||||
for x in reversed(tsink_toposort):
|
||||
if x.op in {Ops.DEVICE, Ops.UNIQUE}: continue
|
||||
|
||||
# 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
|
||||
|
||||
if x.dtype.scalar() == dtypes.weakint: continue # TODO: why do I need this?
|
||||
ending_ranges[x] = sum([ending_ranges.get(u, []) for u in consumer_map[x]], [])
|
||||
|
||||
# *** 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 = [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(dtypes.bool, False).usum(valids)
|
||||
_out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), 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)])
|
||||
|
||||
# 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)
|
||||
# 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(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys())
|
||||
|
||||
# REDUCE creates ranges for the axes it is reducing
|
||||
if x.op is Ops.REDUCE and len(x.arg[1]):
|
||||
if IMAGE:
|
||||
reduce_body = reduce_bodies[x]
|
||||
if any(other.arg[1] != x.arg[1] and x not in body and other not in reduce_body and
|
||||
any(u.op not in ALWAYS_CONTIGUOUS for u in reduce_body & body)
|
||||
for other,body in reduce_bodies.items() if other is not x):
|
||||
rctx.realize_map[x.src[0]] = None
|
||||
rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
|
||||
|
||||
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")
|
||||
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)
|
||||
Reference in New Issue
Block a user