IQ.Pilot Prebuilt Release @ 7e87bc7
This commit is contained in:
IQ.Lvbs CI [bot]
203
artifacts/package_runtime/tinygrad/schedule/__init__.py
Normal file
203
artifacts/package_runtime/tinygrad/schedule/__init__.py
Normal file
@@ -0,0 +1,203 @@
|
||||
import time, inspect
|
||||
from collections import deque
|
||||
from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, rewrite_group, graph_rewrite, gate_kernel_sink, KernelInfo
|
||||
from tinygrad.uop.spec import type_verify, spec_tensor
|
||||
from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup
|
||||
|
||||
# **** schedule linearizer
|
||||
|
||||
# unwrap VIEW/CAST/etc to find the actual data source (kernel output, buffer, or multi-device op)
|
||||
def _unwrap_src(s: UOp) -> UOp:
|
||||
while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}: s = s.src[0]
|
||||
return s
|
||||
|
||||
# a buffer state is AFTER | BUFFER | PARAM. MSELECT/MSTACK join per-device states, BIND is not a buffer dependency
|
||||
def _states(s: UOp) -> list[UOp]:
|
||||
s = _unwrap_src(s)
|
||||
if s.op in {Ops.MSELECT, Ops.MSTACK}: return [st for ss in s.src for st in _states(ss)]
|
||||
if s.op is Ops.BIND: return []
|
||||
assert s.op in {Ops.AFTER, Ops.BUFFER, Ops.PARAM}, f"input to kernel must resolve to a buffer state, not {s.op}"
|
||||
return [s]
|
||||
|
||||
def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]:
|
||||
kernels, remaining = partition(after.src[1:], lambda s: s.op in {Ops.CALL, Ops.END})
|
||||
deps, remaining = partition(remaining, lambda s: s.op is Ops.AFTER)
|
||||
if invalid := [s for s in remaining if s.op is not Ops.STORE]:
|
||||
raise AssertionError(f"AFTER source should be CALL, END, STORE, or AFTER, not {invalid[0].op}")
|
||||
return tuple(kernels), tuple(deps)
|
||||
|
||||
def create_schedule(sched_sink:UOp) -> UOp:
|
||||
with cpu_profile(TracingKey("toposort sched_sink")):
|
||||
# build kernel dependency graph: edges from producer kernel to consumer kernels
|
||||
children: dict[UOp, list[UOp]] = {}
|
||||
in_degree: dict[UOp, int] = {}
|
||||
writes: dict[UOp, list[tuple[UOp, UOp, tuple[UOp, ...]]]] = {} # buffer -> (AFTER, prior state, new kernels)
|
||||
reads: list[tuple[UOp, UOp, UOp]] = [] # (reader AFTER, reader kernel, buffer state read)
|
||||
for u in sched_sink.toposort(gate_kernel_sink):
|
||||
if u.op is not Ops.AFTER: continue
|
||||
kernels, after_deps = _split_after(u)
|
||||
prev_state = _unwrap_src(u.src[0])
|
||||
prev_kernels = set(_split_after(prev_state)[0]) if prev_state.op is Ops.AFTER else set()
|
||||
writes.setdefault(u.buf_uop, []).append((u, prev_state, tuple(k for k in kernels if k not in prev_kernels)))
|
||||
for k in kernels:
|
||||
in_degree.setdefault(k, 0)
|
||||
if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
|
||||
kernel_deps = k.src[0].src[1:] if k.op is Ops.END else k.src[1:]
|
||||
read_states = [st for s in kernel_deps for st in _states(s)]
|
||||
reads += [(u, k, st) for st in read_states]
|
||||
# RAW deps: a kernel runs after the kernels that produced the states it reads or joins
|
||||
for st in read_states + [st for s in after_deps for st in _states(s)]:
|
||||
if st.op is Ops.AFTER:
|
||||
for t in _split_after(st)[0]:
|
||||
children.setdefault(t, []).append(k)
|
||||
in_degree[k] += 1
|
||||
# WAR deps: a kernel reading buffer state S must run before another write that supersedes S. an AFTER only
|
||||
# supersedes its immediate prior state; join members already present in that prior state are ordering deps, not writes
|
||||
for u, k, s in reads:
|
||||
for a, prev_state, write_kernels in writes.get(s.buf_uop, []):
|
||||
if a is u or prev_state is not s: continue
|
||||
for t in write_kernels:
|
||||
if t is not k and t not in k.backward_slice:
|
||||
children.setdefault(k, []).append(t)
|
||||
in_degree[t] += 1
|
||||
|
||||
with cpu_profile(TracingKey("linearize schedule")):
|
||||
queue: deque[UOp] = deque(k for k,v in in_degree.items() if v == 0)
|
||||
linearized: list[UOp] = []
|
||||
while len(queue):
|
||||
rk = queue.popleft()
|
||||
if rk.op is Ops.LINEAR:
|
||||
linearized.extend(rk.src)
|
||||
else:
|
||||
k = rk.src[0] if rk.op is Ops.END else rk
|
||||
assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
|
||||
buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
|
||||
linearized.append(k.src[0].call(*buf_uops))
|
||||
for x in children.get(rk, []):
|
||||
in_degree[x] -= 1
|
||||
if in_degree[x] == 0: queue.append(x)
|
||||
if any(in_degree.values()): raise RuntimeError("cycle detected in assign graph")
|
||||
return UOp(Ops.LINEAR, src=tuple(linearized))
|
||||
|
||||
from tinygrad.schedule.memory import memory_plan_rewrite
|
||||
from tinygrad.engine.realize import capturing, pm_flatten_linear
|
||||
from tinygrad.schedule.rangeify import get_kernel_graph
|
||||
from tinygrad.helpers import CAPTURING
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
|
||||
from tinygrad.dtype import AddrSpace
|
||||
|
||||
def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
|
||||
if (ret:=ctx[0].get(b, None)) is None: ctx[0][b] = ret = UOp.new_buffer(b.device, b.max_numel(), b.dtype)
|
||||
return ret
|
||||
|
||||
pm_post_sched_cache = PatternMatcher([
|
||||
# only resolve buffer PARAMs (slot>=0); ALU/shape vars use slot=-1 and must not be swapped for call args
|
||||
(UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg.slot] if x.arg.slot >= 0 else None),
|
||||
# create new BUFFERs
|
||||
(UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
|
||||
create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
|
||||
])
|
||||
|
||||
def resolve_linear_call(linear_call:UOp):
|
||||
linear = graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")
|
||||
binds = {f"p{i}":x.src[0] for i,x in enumerate(linear_call.src[1:]) if x.op is Ops.BIND}
|
||||
return linear.substitute({v:binds[v.expr] for v in linear.variables() if v.expr in binds}, enter_calls=True, name="resolve scalar params")
|
||||
|
||||
pm_resolve_linear_call = PatternMatcher([
|
||||
# call LINEAR is resolved here
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), resolve_linear_call),
|
||||
])+pm_flatten_linear
|
||||
|
||||
schedule_cache: dict[bytes, UOp] = {}
|
||||
# ctx is just for DEBUG on inner
|
||||
def lower_sink_to_linear(function:UOp) -> UOp|None:
|
||||
st = time.perf_counter()
|
||||
if isinstance(function.arg, KernelInfo): return None
|
||||
cache_key = function.key
|
||||
if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
|
||||
if SPEC: type_verify(function, spec_tensor)
|
||||
# support recursive CALLs
|
||||
linear = create_schedule(get_kernel_graph(function))
|
||||
if SCACHE: schedule_cache[cache_key] = linear
|
||||
else:
|
||||
# schedule cache hit
|
||||
linear = sc_ret
|
||||
if (DEBUG >= 1 and len(linear.src) > 1) or DEBUG >= 3:
|
||||
for frm in inspect.stack():
|
||||
if frm.filename == "<string>": continue
|
||||
if frm.filename.startswith(str(BASEDIR / "apps")): break
|
||||
if not frm.filename.startswith(str(BASEDIR)) and not frm.filename.endswith("/contextlib.py"): break
|
||||
else:
|
||||
frm = None
|
||||
print(f"scheduled {len(linear.src):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
|
||||
f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {cache_key.hex()[:8]}"+\
|
||||
f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
|
||||
return linear
|
||||
|
||||
pm_schedule = PatternMatcher([
|
||||
(UPat(Ops.SINK, name="function"), lower_sink_to_linear),
|
||||
])
|
||||
|
||||
def assert_all_same_devices(ast:UOp):
|
||||
devices = dedup([x.device for x in ast.toposort() if x.op is Ops.PARAM and x.device is not None])
|
||||
if len(devices) >= 2: raise RuntimeError(f"all buffers must be on the same device: {devices}")
|
||||
|
||||
def copy_kernel_to_copy_uop(call:UOp, dst:UOp, src:UOp, r:UOp|None=None):
|
||||
if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None
|
||||
return call.replace(src=(UOp(Ops.COPY, src=(src,), arg=dst.device),) + call.src[1:])
|
||||
|
||||
def simplify_copy_kernel(call:UOp, ast:UOp, dst:UOp, src:UOp):
|
||||
# NOTE: this is a codegen for SDMA devices
|
||||
if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None
|
||||
from tinygrad.codegen.simplify import pm_flatten_range, pm_simplify_ranges
|
||||
from tinygrad.schedule.rangeify import pm_mops
|
||||
from tinygrad.uop.symbolic import sym
|
||||
sink = graph_rewrite(ast, sym+pm_mops+pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges in copy")
|
||||
return call.replace(src=(sink,) + call.src[1:])
|
||||
|
||||
pm_copy_from_store = PatternMatcher([
|
||||
# simplify copy kernels
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"), UPat.var("dst"), UPat.var("src")), name="call"), simplify_copy_kernel),
|
||||
|
||||
# replace this with a copy if it's a copy
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.CONST, arg=0))
|
||||
.store(UPat(Ops.PARAM, name="src").index(UPat(Ops.CONST, arg=0))).sink(),),
|
||||
name="call", allow_any_len=True), copy_kernel_to_copy_uop),
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.RANGE, name="r"))
|
||||
.store(UPat(Ops.PARAM, name="src").index(UPat(Ops.RANGE, name="r"))).end(UPat(Ops.RANGE, name="r")).sink(),),
|
||||
name="call", allow_any_len=True), copy_kernel_to_copy_uop),
|
||||
|
||||
# if it wasn't copy, it currently can't be cross device
|
||||
(UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices),
|
||||
])
|
||||
|
||||
@rewrite_group(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
|
||||
def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
|
||||
# big_sink srcs are all the Tensors
|
||||
linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
|
||||
|
||||
# this recursively resolves the linear_call and allocates buffers
|
||||
linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")
|
||||
|
||||
# create copies
|
||||
linear = graph_rewrite(linear, pm_copy_from_store, name="create COPY kernels for SDMA")
|
||||
|
||||
# vars used in the schedule
|
||||
used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
|
||||
# get var_vals
|
||||
var_vals: dict[str, int] = {}
|
||||
for b in big_sink.src[1:]:
|
||||
if b.op is Ops.BIND:
|
||||
nm = b.src[0].expr
|
||||
if nm not in used_vars: continue
|
||||
val = b.src[1].val
|
||||
if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
|
||||
var_vals[nm] = val
|
||||
|
||||
# jit captures this schedule, no need to execute.
|
||||
if len(capturing) and CAPTURING:
|
||||
capturing[0].add_linear(linear, var_vals)
|
||||
return UOp(Ops.LINEAR, src=()), var_vals
|
||||
|
||||
held_bufs = ({b for b in linear_call.src[1:] if b.op is Ops.BUFFER} if linear_call.op is Ops.CALL else set())
|
||||
return memory_plan_rewrite(linear, held_bufs), var_vals
|
||||
66
artifacts/package_runtime/tinygrad/schedule/allreduce.py
Normal file
66
artifacts/package_runtime/tinygrad/schedule/allreduce.py
Normal file
@@ -0,0 +1,66 @@
|
||||
import functools, itertools
|
||||
from tinygrad.helpers import all_int, prod, DEBUG, RING, ALL2ALL, getenv
|
||||
from tinygrad.uop.ops import UOp
|
||||
|
||||
# *** allreduce implementation ***
|
||||
def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
|
||||
if not isinstance(buf.device, tuple): return None
|
||||
ndev, shape, numel = len(buf.device), buf.shape, prod(buf.shape)
|
||||
op, device = red.arg
|
||||
|
||||
# ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically)
|
||||
# fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks.
|
||||
concrete = all_int(shape)
|
||||
use_all2all = concrete and (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1))
|
||||
use_ring = concrete and not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
|
||||
if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}")
|
||||
|
||||
if not concrete: buf = buf.pad_to(buf.max_shape)
|
||||
# contiguous before we copy it
|
||||
buf = buf.contiguous()
|
||||
|
||||
# naive: copy to all devices. if you shrink later, that'll be handled
|
||||
if not use_ring and not use_all2all:
|
||||
out = functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)])
|
||||
return out if concrete else out.shrink_to(shape)
|
||||
|
||||
# chunk data into ndev pieces
|
||||
assert isinstance(numel, int)
|
||||
factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1)
|
||||
base, left = divmod(numel // factor, ndev)
|
||||
chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (ndev - left), initial=0)))
|
||||
|
||||
# reduce-scatter
|
||||
reduced_chunks:list[UOp] = []
|
||||
for i,(s,e) in enumerate(chunks):
|
||||
if use_all2all:
|
||||
chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(ndev)]
|
||||
reduced_chunks.append(functools.reduce(lambda x,y: x.alu(op, y), chunks_on_i))
|
||||
else:
|
||||
chunk, reduced = buf.reshape((numel,)).shrink(((s,e),)), buf.reshape((numel,)).shrink(((s,e),))
|
||||
for step in range(ndev-1):
|
||||
src, dest = (i+step)%ndev, (i+step+1)%ndev
|
||||
cp = reduced.copy_to_device(buf.device[dest], src if isinstance(reduced.device, tuple) else None)
|
||||
reduced = cp.alu(op, chunk.copy_to_device(buf.device[dest], dest))
|
||||
reduced_chunks.append(reduced)
|
||||
|
||||
# allgather
|
||||
copied_chunks:list[UOp] = []
|
||||
for i,rc in enumerate(reduced_chunks):
|
||||
if isinstance(device, str): copied_chunks.append(rc.copy_to_device(device))
|
||||
elif use_all2all: copied_chunks.append(UOp.mstack(*(rc.copy_to_device(buf.device[j]) for j in range(ndev))))
|
||||
else:
|
||||
chain:list[UOp] = [rc]
|
||||
for step in range(ndev-1):
|
||||
chain.append(rc := rc.copy_to_device(buf.device[(i+step)%ndev]))
|
||||
copied_chunks.append(UOp.mstack(*(chain[(j-i+1)%ndev] for j in range(ndev))))
|
||||
|
||||
# reassemble
|
||||
return UOp.usum(*[c.pad(((s,numel-e),)) for (s,e),c in zip(chunks, copied_chunks)]).reshape(shape)
|
||||
|
||||
def create_allreduce_function(buf:UOp, red:UOp, output:UOp|None=None) -> UOp|None:
|
||||
if output is None: output = UOp.invalids(red.shape, dtype=red.dtype, device=red.device)
|
||||
to = red.param_like(0)
|
||||
src = buf.param_like(1)
|
||||
red = src.allreduce(*red.arg)
|
||||
return output.after(to.after(to.store(handle_allreduce(src, red))).sink().call(output, buf.contiguous(), name="allreduce", precompile=True))
|
||||
334
artifacts/package_runtime/tinygrad/schedule/indexing.py
Normal file
334
artifacts/package_runtime/tinygrad/schedule/indexing.py
Normal file
@@ -0,0 +1,334 @@
|
||||
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)
|
||||
64
artifacts/package_runtime/tinygrad/schedule/memory.py
Normal file
64
artifacts/package_runtime/tinygrad/schedule/memory.py
Normal file
@@ -0,0 +1,64 @@
|
||||
from collections import defaultdict
|
||||
from tinygrad.helpers import NO_MEMORY_PLANNER, DEBUG, round_up
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.runtime.support.memory import TLSFAllocator
|
||||
|
||||
def _collect_bufs(u:UOp) -> list[UOp]:
|
||||
if u.op is Ops.BUFFER: return [u]
|
||||
if u.op in {Ops.MSELECT, Ops.MSTACK}: return [b for s in u.src for b in _collect_bufs(s)]
|
||||
return []
|
||||
|
||||
def _can_plan(b:UOp, held_bufs:set[UOp]) -> bool:
|
||||
if b in held_bufs: return False
|
||||
devs = (b.device,) if isinstance(b.device, str) else b.device
|
||||
# CL and WEBGPU do not support views, see explanation in contiguous_view_offset
|
||||
return all(not d.startswith(("DISK", "TINYFS", "CL", "WEBGPU")) for d in devs)
|
||||
|
||||
LaneKey = tuple[str, int]
|
||||
|
||||
def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp:
|
||||
if NO_MEMORY_PLANNER: return linear
|
||||
if held_bufs is None: held_bufs = set()
|
||||
|
||||
# compute lifetimes for all plannable internal buffers
|
||||
first_appearance:dict[UOp, int] = {}
|
||||
last_appearance:dict[UOp, int] = {}
|
||||
copy_bufs: set[UOp] = set()
|
||||
for i, si in enumerate(linear.src):
|
||||
si_bufs = [b for src in si.src[1:] for b in _collect_bufs(src) if _can_plan(b, held_bufs)]
|
||||
for b in si_bufs:
|
||||
if b not in first_appearance: first_appearance[b] = i
|
||||
last_appearance[b] = i
|
||||
if si.src[0].op is Ops.COPY: copy_bufs.update(si_bufs)
|
||||
if not first_appearance: return linear
|
||||
|
||||
# separate copy and compute buffers into different lanes to avoid introducing dependencies (copy->compute->copy)
|
||||
def _key(b:UOp): return (b.device, 1 if b in copy_bufs else 0)
|
||||
buf_hold = {b: last_appearance[b] - first_appearance[b] + 1 for b in first_appearance if b in copy_bufs}
|
||||
|
||||
# suballocation: build sorted open/close events, then alloc/free in order
|
||||
block_size = 256
|
||||
nbytes = {b: round_up(b.max_numel() * b.dtype.itemsize, block_size) for b in first_appearance}
|
||||
events = sorted([(first_appearance[b], True, b) for b in first_appearance] +
|
||||
[(last_appearance[b] + 1 + buf_hold.get(b, 0), False, b) for b in first_appearance], key=lambda x: (x[0], x[1]))
|
||||
total_memory = sum(nbytes.values()) * 2
|
||||
|
||||
offsets:dict[UOp, int] = {}
|
||||
peaks:dict[LaneKey, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=block_size, lv2_cnt=32)))
|
||||
for _, is_open, buf in events:
|
||||
if is_open: offsets[buf] = peaks[_key(buf)][1].alloc(nbytes[buf])
|
||||
else: peaks[_key(buf)][1].free(offsets[buf])
|
||||
peaks[_key(buf)] = (max(peaks[_key(buf)][0], offsets[buf] + buf.max_numel() * buf.dtype.itemsize), peaks[_key(buf)][1])
|
||||
arena_sizes = {key: round_up(peak, block_size) for key, (peak, _) in peaks.items()}
|
||||
|
||||
# build replace_map: each buffer becomes a SLICE into a shared per-device-lane arena
|
||||
arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()}
|
||||
replace_map:dict[UOp, UOp] = {}
|
||||
for buf_uop, offset in offsets.items():
|
||||
replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(offset)), buf_uop.max_numel())
|
||||
|
||||
if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6):
|
||||
print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs")
|
||||
|
||||
return linear.substitute(replace_map, name="memory plan", walk=True)
|
||||
319
artifacts/package_runtime/tinygrad/schedule/multi.py
Normal file
319
artifacts/package_runtime/tinygrad/schedule/multi.py
Normal file
@@ -0,0 +1,319 @@
|
||||
from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
|
||||
from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, AxisType, graph_rewrite, broadcast_axes, _broadcast_shape, sint_to_uop
|
||||
from tinygrad.uop.ops import sint, ssimplify
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.schedule.allreduce import handle_allreduce
|
||||
|
||||
# ***** multi rewrite MSELECT/MSTACK *****
|
||||
|
||||
def _apply_shrink(marg, s:UOp, i:int) -> UOp:
|
||||
new_arg = [tuple([x.substitute({drng[0]:drng[0].const_like(i)}) if isinstance(x, UOp) and
|
||||
(drng:=[r for r in x.ranges if r.arg[-1] is AxisType.DEVICE]) else x for x in ss]) for ss in marg]
|
||||
return s._mop(Ops.SHRINK, tuple(new_arg))
|
||||
|
||||
def mstack_early_shrink(ms:UOp, shrink:UOp):
|
||||
ret:list[UOp] = []
|
||||
for i, x in enumerate(ms.src):
|
||||
if x.op is Ops.COPY:
|
||||
ret.append(_apply_shrink(shrink.marg, x.src[0], i).copy_to_device(x.device))
|
||||
else:
|
||||
ret.append(_apply_shrink(shrink.marg, x, i).contiguous())
|
||||
return ms.replace(src=tuple(ret))
|
||||
|
||||
def lower_broadcast_copy(c:UOp, x:UOp):
|
||||
if not (isinstance(c.device, tuple) and isinstance(x.device, str)): return None
|
||||
if (sx:=x.simplify()).device is None and sx.base.op is Ops.CONST: return UOp(Ops.MSTACK, src=(sx,)*len(c.device))
|
||||
return UOp(Ops.MSTACK, src=tuple(x.copy_to_device(d) for d in c.device))
|
||||
|
||||
replace_allreduce = PatternMatcher([
|
||||
# BROADCAST: explicitly expand broadcast copies and combine with MSTACK
|
||||
(UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lower_broadcast_copy),
|
||||
# COPY_TO_ONE: if copying from multidevice to one, MSELECT the first (TODO: a little from each?)
|
||||
(UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lambda c,x:
|
||||
x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None),
|
||||
# MSELECT on MSTACK is replaced with nothing
|
||||
(UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),
|
||||
# move shrink before MSTACK
|
||||
(UPat(Ops.SHRINK, src=(UPat(Ops.MSTACK, name="ms"),), allow_any_len=True, name="shrink"), mstack_early_shrink),
|
||||
# move MSELECT before movement ops
|
||||
(UPat(Ops.MSELECT, src=(UPat(GroupOp.Movement, src=(UPat.var("s"),), allow_any_len=True, name="v"),), name="ms"),
|
||||
lambda s,v,ms: v.replace(src=(s.mselect(ms.arg),)+v.src[1:])),
|
||||
])
|
||||
|
||||
_early_allreduce = PatternMatcher([
|
||||
(UPat(Ops.ALLREDUCE, src=(UPat.var("buf"),), name="red"), handle_allreduce),
|
||||
])
|
||||
if not getenv("LATE_ALLREDUCE", 1): replace_allreduce = _early_allreduce + replace_allreduce
|
||||
|
||||
# ***** multi functions *****
|
||||
|
||||
def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
|
||||
# normalize srcs to local shards on axis
|
||||
devices = [x.device for x in msrcs if x.device is not None]
|
||||
assert all_same(devices), f"all buffers must have the same device {devices}"
|
||||
# without devices the sharding range comes from the UNSHARD itself (e.g. a LOCAL thread range);
|
||||
# device shards range over the devices instead
|
||||
if len(devices): sharding_rng = UOp.range(len(devices[0]), -1, AxisType.DEVICE)
|
||||
else:
|
||||
sharding_rng = next((m.src[1] for m in msrcs if m.op is Ops.UNSHARD), None)
|
||||
assert sharding_rng is not None, "shard_srcs requires a device or a sharding range"
|
||||
|
||||
out_shape = _broadcast_shape(*[x.shape for x in msrcs])
|
||||
srcs:list[UOp] = []
|
||||
for mlb in msrcs:
|
||||
src_axis = axis - (len(out_shape)-len(mlb.shape))
|
||||
if mlb.axis == src_axis:
|
||||
# same axis, just copy through
|
||||
srcs.append(mlb.src[0])
|
||||
else:
|
||||
# otherwise every shard gets the full copy, sharded iff this src has the axis (broadcast srcs stay whole)
|
||||
full = mlb if mlb.axis is None else copy_multi(mlb, mlb.device)
|
||||
srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, sharding_rng))
|
||||
return srcs
|
||||
|
||||
def shard_subview(full:UOp, multi:UOp) -> UOp:
|
||||
"""the sub-view of an unsharded full-shape value (shape == multi.shape) that belongs to this shard:
|
||||
_shard along every sharded axis (contiguous blocks, like the device path)."""
|
||||
assert tuple(full.shape) == tuple(multi.shape), f"shard sub-view shape mismatch {full.shape} != {multi.shape}"
|
||||
# an EXPAND of a scalar over the full shape is the same broadcast on every shard: re-expand over the shard shape
|
||||
if full.op is Ops.EXPAND and full.src[0].shape == (): return full.src[0].expand(multi.src[0].shape)
|
||||
for ax, rng in multi.sharding: full = full._shard(ax, rng)
|
||||
return full
|
||||
|
||||
def alu_multi(root:UOp):
|
||||
multis = [m for m in root.src if m.op is Ops.UNSHARD]
|
||||
if not multis: return None
|
||||
sharding = multis[0].sharding
|
||||
target = multis[0]
|
||||
def can_handle(m:UOp) -> bool:
|
||||
# same sharding (peel the UNSHARD), or a whole unsharded value of the full tile shape (takes its per-shard
|
||||
# sub-view), or a broadcast scalar
|
||||
if m.sharding: return m.sharding == sharding
|
||||
return m.shape == () or tuple(m.shape) == tuple(target.shape)
|
||||
if all(can_handle(m) for m in root.src):
|
||||
# every src either has the target sharding or is whole on every shard: run the alu per-shard
|
||||
srcs = [m.src[0] if m.op is Ops.UNSHARD else m if m.shape == () else shard_subview(m, target) for m in root.src]
|
||||
return srcs[0].alu(root.op, *srcs[1:]).unshard(target.arg, target.src[1:])
|
||||
# resharding: single-axis fallback via shard_srcs
|
||||
axis = root.axis
|
||||
assert axis is not None
|
||||
srcs = shard_srcs(root.src, axis)
|
||||
return srcs[0].alu(root.op, *srcs[1:]).unshard(axis, next(m.src[1] for m in root.src if m.op is Ops.UNSHARD))
|
||||
|
||||
def reduce_multi(root:UOp, multi:UOp):
|
||||
op, num_axes = root.arg
|
||||
sharding = multi.sharding
|
||||
reduced = [(ax, rng) for ax, rng in sharding if ax < num_axes]
|
||||
remaining = [(ax, rng) for ax, rng in sharding if ax >= num_axes]
|
||||
local = multi.src[0]._rop(op, tuple(range(num_axes)))
|
||||
if reduced:
|
||||
assert not remaining, f"partial allreduce not supported for multi-axis sharding {sharding}"
|
||||
# all sharded axes are reduced: full allreduce
|
||||
if ALLREDUCE_CAST and multi.src[0].op is Ops.CAST and multi.src[0].src[0].dtype in (dtypes.bfloat16, dtypes.half):
|
||||
orig_dtype = multi.src[0].src[0].dtype
|
||||
return local.cast(orig_dtype).allreduce(op, multi.device).cast(local.dtype)
|
||||
return local.allreduce(op, multi.device)
|
||||
# no sharded axes reduced: piecewise, keep all remaining sharding
|
||||
new_axes = tuple(ax - num_axes for ax, _ in remaining)
|
||||
new_rngs = tuple(rng for _, rng in remaining)
|
||||
return local.unshard(new_axes, new_rngs)
|
||||
|
||||
def reshape_multi(root:UOp, multi:UOp):
|
||||
if prod(multi.shape) != prod(new_shape:=root.marg): raise RuntimeError("reshape must maintain prod(shape)")
|
||||
# map every sharded axis through the reshape: the axis boundary must survive intact and stay divisible by its shard count
|
||||
arg_acc:list[sint] = [1]
|
||||
for s in new_shape: arg_acc.append(ssimplify(arg_acc[-1]*s))
|
||||
new_shardings = []
|
||||
for ax, rng in multi.sharding:
|
||||
count = int(rng.vmax)+1
|
||||
target = prod(multi.shape[:ax])
|
||||
if target not in arg_acc: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
|
||||
new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1
|
||||
if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
|
||||
new_shardings.append((new_ax, rng))
|
||||
new_axs = {a for a, _ in new_shardings}
|
||||
new_shape = tuple(s//(int(rng.vmax)+1) if a in new_axs else s for a,s in enumerate(new_shape))
|
||||
return multi.src[0].reshape(new_shape).unshard(tuple(a for a,_ in new_shardings), tuple(r for _,r in new_shardings))
|
||||
|
||||
def expand_multi(root:UOp, multi:UOp):
|
||||
shift = len(root.marg)
|
||||
return multi.src[0]._mop(Ops.EXPAND, arg=root.marg) \
|
||||
.unshard(tuple(ax+shift for ax,_ in multi.sharding), tuple(r for _,r in multi.sharding))
|
||||
|
||||
def pad_multi(root:UOp, multi:UOp):
|
||||
for ax, _ in multi.sharding:
|
||||
assert root.marg[ax] == (0, multi.shape[ax]), f"padding not supported for {root.marg=}"
|
||||
counts = {a for a,_ in multi.sharding}
|
||||
local_pad = tuple((0, multi.src[0].shape[a]) if a in counts else s for a,s in enumerate(root.marg))
|
||||
return multi.src[0]._mop(Ops.PAD, local_pad).unshard(multi.arg, multi.src[1:])
|
||||
|
||||
def permute_multi(root:UOp, multi:UOp):
|
||||
# all permutes supported!
|
||||
return multi.src[0].permute(root.marg) \
|
||||
.unshard(tuple(root.marg.index(ax) for ax,_ in multi.sharding), tuple(r for _,r in multi.sharding))
|
||||
|
||||
def shrink_multi(root:UOp, multi:UOp):
|
||||
# resolve each sharded axis independently: a shrink to exactly this range's own shard resolves the UNSHARD along
|
||||
# that axis (e.g. a fragment indexed by its LOCAL thread range becomes that thread's REG shard, no copy needed)
|
||||
local_marg = list(root.marg)
|
||||
remaining = list(multi.sharding)
|
||||
for ax, rng in multi.sharding:
|
||||
shard_sz = multi.src[0].shape[ax]
|
||||
s, l = root.marg[ax] # SHRINK marg is (start, length)
|
||||
if sint_to_uop(l).ssimplify() == shard_sz and (sint_to_uop(s)-rng*shard_sz).ssimplify() == 0:
|
||||
local_marg[ax] = (0, shard_sz)
|
||||
remaining.remove((ax, rng))
|
||||
continue
|
||||
part_bounds = tuple((i*shard_sz, shard_sz) for i in range(int(rng.vmax)+1))
|
||||
if (s, l) == (0, multi.shape[ax]): local_marg[ax] = (0, shard_sz) # full axis stays sharded, shrink the other axes locally
|
||||
else:
|
||||
# NOTE: otherwise a shrink on the shard axis is only allowed on the legacy device path, selecting a single
|
||||
# partition (which is copied to all the devices and optimized out later)
|
||||
if len(multi.sharding) != 1 or not isinstance(multi.device, tuple) or (s, l) not in part_bounds:
|
||||
raise RuntimeError(f"shrinking not supported for {root.marg=}")
|
||||
non_shard_shrink = tuple((0, shard_sz) if i == ax else t for i, t in enumerate(root.marg))
|
||||
return multi.src[0].copy_to_device(multi.device, arg=part_bounds.index((s, l)))._mop(Ops.SHRINK, non_shard_shrink)
|
||||
val = multi.src[0]._mop(Ops.SHRINK, tuple(local_marg))
|
||||
return val if not remaining else val.unshard(tuple(a for a,_ in remaining), tuple(r for _,r in remaining))
|
||||
|
||||
def flip_multi(root:UOp, multi:UOp):
|
||||
for ax, _ in multi.sharding:
|
||||
if root.marg[ax]: raise RuntimeError(f"flipping not supported on sharded axis {ax}")
|
||||
return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).unshard(multi.arg, multi.src[1:])
|
||||
|
||||
def stack_multi(root:UOp):
|
||||
# STACK adds a leading axis: srcs are sharded one axis below the output
|
||||
multis = [m for m in root.src if m.op is Ops.UNSHARD]
|
||||
if not multis: return None
|
||||
sharding = multis[0].sharding
|
||||
if all(m.sharding == sharding for m in multis):
|
||||
srcs = [m.src[0] if m.op is Ops.UNSHARD else m for m in root.src]
|
||||
new_sharding = tuple((ax+1, rng) for ax, rng in sharding)
|
||||
return UOp(Ops.STACK, src=tuple(srcs)).unshard(tuple(a for a,_ in new_sharding), tuple(r for _,r in new_sharding))
|
||||
# resharding: single-axis fallback
|
||||
axis = root.axis
|
||||
assert axis is not None
|
||||
return UOp(Ops.STACK, src=tuple(shard_srcs(root.src, axis-1))).unshard(axis, next(m.src[1] for m in root.src if m.op is Ops.UNSHARD))
|
||||
|
||||
def index_multi(root:UOp, multi:UOp):
|
||||
# INDEX on UNSHARD: resolve each sharded axis into this range's own shard.
|
||||
# Two ownership patterns are supported:
|
||||
# contiguous: idx = rng*shard_sz + local (thread rng owns [rng*shard_sz, ...))
|
||||
# strided: idx = rng + ir*shard_sz (thread rng owns {rng, rng+shard_sz, ...})
|
||||
idxs = list(root.src[1:])
|
||||
for ax, rng in multi.sharding:
|
||||
shard_sz = multi.src[0].shape[ax]
|
||||
local = (idxs[ax] - rng*shard_sz).simplify()
|
||||
if local.vmin >= 0 and local.vmax < shard_sz:
|
||||
idxs[ax] = local
|
||||
continue
|
||||
# strided ownership: idx ≡ rng (mod shard_sz), intra-shard position is (idx - rng) // shard_sz
|
||||
diff = (idxs[ax] - rng).simplify()
|
||||
if (mod:=(diff % shard_sz).simplify()).op is Ops.CONST and mod.val == 0:
|
||||
local = (diff // shard_sz).simplify()
|
||||
if local.vmin >= 0 and local.vmax < shard_sz:
|
||||
idxs[ax] = local
|
||||
continue
|
||||
raise RuntimeError(f"index_multi: cannot shard index {idxs[ax]} for UNSHARD axis {ax} with shard size {shard_sz}")
|
||||
return multi.src[0].index(*idxs)
|
||||
|
||||
def _shard_idx(rng:UOp, dev_idx:int) -> int:
|
||||
drngs = [r for r in rng.ranges if r.arg[-1] is AxisType.DEVICE]
|
||||
return 0 if not drngs else int(rng.substitute({drngs[0]: drngs[0].const_like(dev_idx)}).ssimplify())
|
||||
|
||||
def copy_multi(multi:UOp, device:str | tuple[str, ...]):
|
||||
sharding = multi.sharding
|
||||
if isinstance(device, str):
|
||||
# reconstruct by concatenating along each axis from last to first
|
||||
piece_info: list[tuple[tuple, UOp]] = []
|
||||
for i in range(len(multi.device)):
|
||||
idxs = tuple(_shard_idx(r, i) for _, r in sharding)
|
||||
piece_info.append((idxs, multi.src[0].mselect(i).copy_to_device(device)))
|
||||
for j in range(len(sharding) - 1, -1, -1):
|
||||
ax, rng = sharding[j]
|
||||
groups: dict[tuple, list[tuple[int, UOp]]] = {}
|
||||
for idxs, p in piece_info:
|
||||
key = idxs[:j] + idxs[j+1:]
|
||||
groups.setdefault(key, []).append((idxs[j], p))
|
||||
piece_info = []
|
||||
for key in sorted(groups):
|
||||
grp = sorted(groups[key], key=lambda x: x[0])
|
||||
piece_info.append((key, grp[0][1].cat(*[x[1] for x in grp[1:]], dim=ax)))
|
||||
return piece_info[0][1]
|
||||
# multi-device target: unshard all axes and allreduce
|
||||
val = multi.src[0]
|
||||
for ax, rng in sharding:
|
||||
bsz = val.shape[ax]
|
||||
val = val.pad(tuple((0,0) if a != ax else (bsz*rng, bsz*int(rng.vmax) - bsz*rng) for a in range(len(val.shape))))
|
||||
return val.allreduce(Ops.ADD, device)
|
||||
|
||||
def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).unshard(src.arg, src.src[1:])
|
||||
|
||||
def store_value_multi(dest:UOp, multi:UOp):
|
||||
# storing a sharded value into an unsharded dest: every shard stores into its own sub-view of the dest
|
||||
return shard_subview(dest, multi).store(multi.src[0])
|
||||
|
||||
def store_dest_multi(root:UOp, multi:UOp):
|
||||
# STORE with a sharded dest: every shard stores into its own shard of the dest.
|
||||
# the value is handled like in alu_multi: UNSHARD srcs peel, full-shape values take their per-shard sub-view
|
||||
# (scalars arrive EXPANDed to the full shape by UOp.store's const_like, so they sub-view like everything else)
|
||||
srcs = [multi.src[0]] + [x.src[0] if x.op is Ops.UNSHARD else shard_subview(x, multi) if tuple(x.shape) == tuple(multi.shape) else x
|
||||
for x in root.src[1:]]
|
||||
return UOp(root.op, root.dtype, tuple(srcs), root.arg)
|
||||
|
||||
def passthrough_multi(root:UOp, multi:UOp):
|
||||
new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:])
|
||||
return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:])
|
||||
|
||||
def rewrite_into_function(call:UOp):
|
||||
if call.arg.precompile: return None
|
||||
new_body = graph_rewrite(call.src[0], multi_pm, name="subcall")
|
||||
new_args = tuple(a.src[0] if a.op is Ops.UNSHARD else a for a in call.src[1:])
|
||||
# after multi resolution, TUPLE elements may be UNSHARD — strip UNSHARD from body, create per-shard FUNCTION, wrap each GETTUPLE in its own UNSHARD
|
||||
assert new_body.op is Ops.TUPLE
|
||||
if any(s.op is Ops.UNSHARD for s in new_body.src):
|
||||
shard_call = call.replace(src=(UOp.maketuple(*[s.src[0] if s.op is Ops.UNSHARD else s for s in new_body.src]),)+new_args)
|
||||
return UOp.maketuple(*[shard_call.gettuple(i).unshard(s.arg, s.src[1:]) if s.op is Ops.UNSHARD else shard_call.gettuple(i)
|
||||
for i, s in enumerate(new_body.src)])
|
||||
return call.replace(src=(new_body,)+new_args)
|
||||
|
||||
def param_to_multi(p:UOp):
|
||||
if p.axis is None: return None
|
||||
return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name, p.arg.addrspace).unshard(p.axis)
|
||||
|
||||
# NOTE: this is the same pattern as unrolled ranges
|
||||
multi_pm = PatternMatcher([
|
||||
(UPat(Ops.PARAM, name="p"), param_to_multi),
|
||||
(UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.UNSHARD])), alu_multi),
|
||||
(UPat(Ops.REDUCE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), reduce_multi),
|
||||
(UPat(Ops.RESHAPE, src=(UPat(Ops.UNSHARD, name="multi"), UPat()), name="root"), reshape_multi),
|
||||
(UPat(Ops.EXPAND, src=(UPat(Ops.UNSHARD, name="multi"), UPat()), name="root"), expand_multi),
|
||||
(UPat(Ops.PAD, src=(UPat(Ops.UNSHARD, name="multi"), UPat(), UPat()), name="root"), pad_multi),
|
||||
(UPat(Ops.SHRINK, src=(UPat(Ops.UNSHARD, name="multi"), UPat(), UPat()), name="root"), shrink_multi),
|
||||
(UPat(Ops.PERMUTE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), permute_multi),
|
||||
(UPat(Ops.FLIP, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), flip_multi),
|
||||
(UPat(Ops.STACK, name="root", custom_early_reject=set([Ops.UNSHARD])), stack_multi),
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.UNSHARD, name="multi"),), name="root", allow_any_len=True), index_multi),
|
||||
(UPat(Ops.AFTER, src=(UPat(Ops.UNSHARD), UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="dest"), UPat(Ops.UNSHARD, name="src"))))), store_after_multi),
|
||||
(UPat(Ops.COPY, src=(UPat(Ops.UNSHARD, name="multi"),), name="copy"), lambda multi,copy: copy_multi(multi, copy.arg)),
|
||||
(UPat(Ops.ALLREDUCE, src=(UPat(Ops.UNSHARD, name="multi"),), name="red"),
|
||||
lambda multi,red: multi.src[0].allreduce(*red.arg).unshard(multi.arg, multi.src[1:])),
|
||||
|
||||
# resolve TUPLE+GETTUPLE (needed in multi)
|
||||
(UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
|
||||
# GETTUPLE on UNSHARD: passthrough UNSHARD (e.g. when FUNCTION was replaced by UNSHARD(GETTUPLE(...)))
|
||||
(UPat(Ops.GETTUPLE, src=(UPat(Ops.UNSHARD, name="multi"),), name="g"),
|
||||
lambda g, multi: multi.src[0].gettuple(g.arg).unshard(multi.arg, multi.src[1:]) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE} else multi),
|
||||
# rewrite into FUNCTION calls explicitly for UNSHARD (value-producing)
|
||||
(UPat(Ops.FUNCTION, name="call"), rewrite_into_function),
|
||||
(UPat((Ops.CALL, Ops.FUNCTION, Ops.AFTER), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
|
||||
# just strip the UNSHARD from non-value-producing CALLs (custom kernels, etc.) — FUNCTION is handled by rewrite_into_function
|
||||
(UPat(Ops.CALL, dtype=dtypes.void, name="root", custom_early_reject=set([Ops.UNSHARD])), lambda root:
|
||||
UOp(root.op, root.dtype, tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), root.arg)),
|
||||
(UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
|
||||
src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), passthrough_multi),
|
||||
# STORE of a sharded value into an unsharded dest (e.g. a fragment into a full output tile)
|
||||
(UPat(Ops.STORE, src=(UPat.var("dest"), UPat(Ops.UNSHARD, name="multi"))), store_value_multi),
|
||||
# STORE into a sharded dest (e.g. the fragment init): every shard stores into its own shard
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), store_dest_multi),
|
||||
])+replace_allreduce
|
||||
579
artifacts/package_runtime/tinygrad/schedule/rangeify.py
Normal file
579
artifacts/package_runtime/tinygrad/schedule/rangeify.py
Normal file
@@ -0,0 +1,579 @@
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import cast
|
||||
import itertools
|
||||
from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
|
||||
from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
|
||||
from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
|
||||
from tinygrad.uop.movement import mop_cleanup
|
||||
from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
|
||||
from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
|
||||
from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
|
||||
from tinygrad.codegen.opt import Opt
|
||||
from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, IndexingContext, apply_movement_op
|
||||
from tinygrad.schedule.multi import multi_pm
|
||||
from tinygrad.schedule.allreduce import create_allreduce_function
|
||||
|
||||
# creation can recurse a lot
|
||||
import sys
|
||||
sys.setrecursionlimit(10000)
|
||||
|
||||
def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
|
||||
if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, after = x.src[0], after.cast(dtypes.float)
|
||||
while True:
|
||||
if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg))
|
||||
elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape)
|
||||
elif x.op is Ops.WHERE and x.src[2].base.is_invalid and x.src[1].op is Ops.PAD:
|
||||
x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (o,_),s in zip(x.src[1].marg, x.src[1].src[0].shape)))
|
||||
else: break
|
||||
ctx[x] = after
|
||||
|
||||
# *** fold moved AFTERs (hack for openpilot) ***
|
||||
pm_fold_moved_after = PatternMatcher([
|
||||
(UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement,Ops.CAST,Ops.WHERE), name="src")))), name="after"), found_after),
|
||||
# replace ALU sources with AFTER versions found above
|
||||
(UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
|
||||
])
|
||||
|
||||
# movement op on INDEX as a PatternMatcher
|
||||
def _mop_index(r:UOp, idx:UOp):
|
||||
idxs = idx.src[1:]
|
||||
if len(idxs) == len(r.shape):
|
||||
return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
|
||||
if r.op is Ops.RESHAPE:
|
||||
src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
|
||||
if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
|
||||
if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
|
||||
ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
|
||||
return ret if ret.shape == idx.shape else None
|
||||
|
||||
pm_mops = PatternMatcher([
|
||||
# handle movement ops on INDEX
|
||||
(UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), _mop_index),
|
||||
# move movement ops and INDEX after AFTER
|
||||
(UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
|
||||
lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
|
||||
(UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
|
||||
])
|
||||
|
||||
# *****************
|
||||
# 0. do some cleanup rewrites, mostly copied from the old stuff
|
||||
|
||||
def fix_store_hazard(target:UOp, src:UOp):
|
||||
if (base:=target.base) not in src.backward_slice_with_self: return None
|
||||
# PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
|
||||
unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
|
||||
reaches_base: dict[UOp, bool] = {}
|
||||
for s in src.toposort(gate=lambda s: s.op is not Ops.CONTIGUOUS):
|
||||
reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src)
|
||||
if reaches_base[s] and s.op in unsafe and not (s is target and s.op is Ops.SHRINK): return target.store(src.contiguous())
|
||||
|
||||
def split_reduceop(reduce:UOp, x:UOp):
|
||||
if prod(reduce.shape) == 0: return None
|
||||
if not SPLIT_REDUCEOP or not all_int(x.shape) or (prod(x.shape)//prod(reduce.shape))<getenv("REDUCEOP_SPLIT_THRESHOLD", 32768): return None
|
||||
# if there are few globals, make some reduces into globals by splitting into two kernels
|
||||
# cap output buffer to 2**22: heuristic number of global outputs to achieve max occupancy with enough locals+upcasts for gemm
|
||||
# ~2**10 should be enough if GROUP is used
|
||||
# 256 split maximum should be "negligible reduce" for low prod(reduce.shape), 8 split minimum.
|
||||
# split is moved to the end to provide maximum locality for the second phase reduce.
|
||||
|
||||
# get expanded by rangeifying the UOp x
|
||||
indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else 0 for i,s in enumerate(x.shape)])
|
||||
range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP, x.base.dtype)}, extra_pm=pm_mops).ranges]
|
||||
is_expanded = [i not in range_nums for i in range(len(x.shape))]
|
||||
|
||||
if not (split_candidates:=[(i,d) for i in range(reduce.arg[1])
|
||||
for d in range(min(256,2**getenv("REDUCEOP_SPLIT_SIZE",22)//prod(reduce.shape)),8-1,-1)
|
||||
if x.shape[i]%d==0 and not is_expanded[i]]): return None
|
||||
dim_to_split, divisor = split_candidates[0]
|
||||
splitted_shape = x.shape[:dim_to_split]+(divisor,)+(x.shape[dim_to_split]//divisor,)+x.shape[dim_to_split+1:]
|
||||
splitted = x.reshape(splitted_shape).permute(tuple([d for d in range(len(splitted_shape)) if d!=dim_to_split]+[dim_to_split]))
|
||||
if DEBUG >= 3: print(f"split {divisor}: {x.shape} -> {splitted.shape} -> {reduce.shape}")
|
||||
# reduce original axes, then split
|
||||
return splitted._rop(reduce.arg[0], tuple(range(reduce.arg[1]))).contiguous()._rop(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape)
|
||||
|
||||
pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p) if p.arg.slot >= 0 else None), ])
|
||||
def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
|
||||
if c.arg.precompile: return None
|
||||
params: list[UOp] = []
|
||||
graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params")
|
||||
params = sorted(params, key=lambda x: x.arg.slot)
|
||||
args = c.src[1:]
|
||||
|
||||
# NOTE: this isn't really needed. it's okay if there's unused args in the function
|
||||
if not allow_param_mismatch:
|
||||
if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}")
|
||||
if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
|
||||
|
||||
dict_map = {x:args[x.arg.slot] for x in params}
|
||||
for i, (p, a) in enumerate(dict_map.items()):
|
||||
if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
|
||||
if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
|
||||
if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
|
||||
return c.src[0].substitute(dict_map, walk=True)
|
||||
|
||||
# shape-changing bitcast
|
||||
def expand_bitcast(bc:UOp) -> UOp|None:
|
||||
x = bc.src[0]
|
||||
if (ns:=bc.dtype.itemsize) == (os:=x.dtype.itemsize) or (isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS"))): return None
|
||||
new_uint, tmp = to_dtype(f"uint{8*ns}"), x.bitcast(to_dtype(f"uint{8*os}"))
|
||||
if ns > os:
|
||||
tmp = tmp.reshape(x.shape[:-1] + (x.shape[-1]//(rate := ns//os), rate))
|
||||
parts = [tmp.shrink((None,)*(len(tmp.shape)-1) + ((i, i+1),)).cast(new_uint)<<8*i*os for i in range(rate)]
|
||||
return parts[0].usum(*parts[1:]).squeeze(-1).bitcast(bc.dtype)
|
||||
parts = [tmp>>8*i*ns for i in range(os//ns)]
|
||||
return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
|
||||
|
||||
earliest_rewrites = mop_cleanup+PatternMatcher([
|
||||
# resolve FUNCTION calls (inline the body)
|
||||
(UPat(Ops.FUNCTION, name="c"), resolve_function),
|
||||
|
||||
# resolve TUPLE+GETTUPLE
|
||||
(UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
|
||||
|
||||
# resolve allreduce (must be bottom up)
|
||||
(UPat(Ops.ALLREDUCE, src=(UPat.var("buf"),), name="red"), create_allreduce_function),
|
||||
|
||||
# split_reduceop
|
||||
(UPat(Ops.REDUCE, name="reduce", src=(UPat.var("x"),)), split_reduceop),
|
||||
|
||||
# remove DETACH/CONTIGUOUS_BACKWARD (TODO: this is copied in allocations)
|
||||
(UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
|
||||
|
||||
# SINK only ever references the base
|
||||
(UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.base for y in x.src))),
|
||||
|
||||
# ** copy rules **
|
||||
|
||||
# COPY transfers a contiguous range, so materialize a source that's resized (shrink/pad/expand) or reordered (permute/flip)
|
||||
(UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"),), name="c"),
|
||||
lambda c,r: c.replace(src=(r.contiguous(),)) if resolve(r.numel() != r.base.numel(), False) or r.contiguous_view_offset() is None else None),
|
||||
|
||||
# copy to same device is a no-op
|
||||
(UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x if x.device == copy.device else None),
|
||||
|
||||
# copy on reshape is reshape on copy
|
||||
(UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
|
||||
|
||||
# reshaping on STORE can be a NOOP
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
|
||||
UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
|
||||
lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
|
||||
|
||||
# ** store rules **
|
||||
|
||||
# fix store hazard (dest is in used in src) by adding contiguous: TestAssign.test_post_flipped_assignment
|
||||
(UPat(Ops.STORE, src=(UPat(name="target"), UPat(name="src"))), fix_store_hazard),
|
||||
|
||||
# remove two STOREs that store the same thing to the same place: TestSchedule.test_dedup_assign
|
||||
(UPat.var("buf").after(UPat.var("buf").store(UPat.var("src")), name="a1").after(UPat.var("a1").store(UPat.var("src"))), lambda buf,src,a1:a1),
|
||||
|
||||
# store a buffer's own current contents back into itself: TestAssign.test_nested_after_contiguous_store_no_init
|
||||
(UPat.var("buf").after(UPat.var("buf").store(UPat.var("buf").after(UPat.var("buf").store(UPat.var("src")), name="a1"))), lambda buf,src,a1:a1),
|
||||
|
||||
# move bitcast from store dest to source: TestAssign.test_assign_bitcast
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))),
|
||||
lambda target, src: target.store(src.bitcast(target.dtype))),
|
||||
|
||||
(UPat(Ops.BITCAST, name="bc"), expand_bitcast),
|
||||
|
||||
# ** size 0 **
|
||||
|
||||
# reduce of size 0 is the identity element
|
||||
(UPat(Ops.REDUCE, name="reduce", src=(UPat.var("x"),)),
|
||||
lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if 0 in x.shape and 0 not in reduce.shape else None),
|
||||
# handle size 0
|
||||
(UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
|
||||
])
|
||||
|
||||
# *****************
|
||||
# 3.5 cleanups
|
||||
|
||||
ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.NOOP}
|
||||
|
||||
# you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
|
||||
def cleanup_dead_axes(b:UOp):
|
||||
if not b.arg.removable: return None
|
||||
# don't optimize ALWAYS_RUN_OPS or AFTER (AFTER is a buffer identity — ranges define consumer access, not computation)
|
||||
if b.src[0].op in ALWAYS_RUN_OPS or b.src[0].op is Ops.AFTER: return None
|
||||
|
||||
new_rng = []
|
||||
hit = False
|
||||
reshape: list[sint] = []
|
||||
for s,rng in zip(b.shape, b.src[1:]):
|
||||
# skip for symbolic. TODO: fix this
|
||||
if rng.op is Ops.RANGE and rng.src[0].op is not Ops.CONST: return None
|
||||
# CONSTs are already dead axes
|
||||
if rng.op is Ops.CONST or (rng.op is Ops.RANGE and rng not in b.src[0].ranges):
|
||||
reshape.append(1)
|
||||
hit = True
|
||||
else:
|
||||
reshape.append(s)
|
||||
new_rng.append(rng)
|
||||
if hit:
|
||||
return b.replace(src=b.src[0:1]+tuple(new_rng)).reshape(tuple(reshape)).expand(b.shape)
|
||||
|
||||
def gate_substitute(ctx, b:UOp) -> None:
|
||||
if not any(r in b.ranges for r in ctx.keys()): raise BottomUpGate()
|
||||
pm_gate_substitute = PatternMatcher([(UPat(GroupOp.All, name="b"), gate_substitute)], compiled=False)
|
||||
# if a buffer is being stored just for permutes or something, remove it
|
||||
# we want to reexpress the indexes of idx2 in terms of the implied b1
|
||||
def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
|
||||
# see if we can't do it, should this ever hit?
|
||||
assert len(buf.src) == len(idx.src), f"index on wrong bufferize, {len(buf.src)} != {len(idx.src)}"
|
||||
assert all(x.op in {Ops.RANGE, Ops.CONST} for x in buf.src[1:])
|
||||
|
||||
# if it's user contiguous, we never remove it
|
||||
if src.op in ALWAYS_RUN_OPS or not buf.arg.removable: return None
|
||||
|
||||
# *** here is where we compute the cost ***
|
||||
# if we return None, the bufferize is kept
|
||||
|
||||
accessed_buffers: list[UOp] = []
|
||||
indexes: list[UOp] = []
|
||||
reduces: list[UOp] = []
|
||||
def red_gate(x:UOp):
|
||||
if x.op is Ops.AFTER:
|
||||
accessed_buffers.append(x.buf_uop)
|
||||
return False
|
||||
if (x.op is Ops.STAGE and x.arg.addrspace == AddrSpace.GLOBAL) or x.op is Ops.MSTACK:
|
||||
accessed_buffers.append(x)
|
||||
return False
|
||||
if x.op is Ops.STORE:
|
||||
# don't look inside stores, this doesn't count toward buffer accesses
|
||||
return False
|
||||
if x.op is Ops.PARAM:
|
||||
accessed_buffers.append(x)
|
||||
if x.op is Ops.INDEX:
|
||||
indexes.append(x)
|
||||
if x.op is Ops.REDUCE: reduces.append(x)
|
||||
return True
|
||||
src.toposort(gate=red_gate)
|
||||
del red_gate
|
||||
accessed_buffers = dedup(accessed_buffers)
|
||||
|
||||
# if this is generated from multiple buffers, don't remove this buffer
|
||||
if len(accessed_buffers) > 3 and not (PCONTIG > 2): return None
|
||||
|
||||
# if any reduces access a buffer, don't remove this buffer
|
||||
buffer_in_reduce = False
|
||||
def buf_gate(x:UOp):
|
||||
nonlocal buffer_in_reduce
|
||||
if x.op in {Ops.PARAM, Ops.STAGE, Ops.AFTER}: buffer_in_reduce = True
|
||||
return not buffer_in_reduce
|
||||
UOp.sink(*[x.src[0] for x in reduces]).toposort(gate=buf_gate)
|
||||
del buf_gate
|
||||
if buffer_in_reduce:
|
||||
if PCONTIG > 2:
|
||||
out_in_ratio = (prod(buf.shape)+1) / (sum([x.numel() for x in accessed_buffers])+1)
|
||||
if out_in_ratio < 10: return None
|
||||
# here we have to check the indexes, we might do a partial contig here
|
||||
local_indexes = [x for x in indexes if x.src[0].op is Ops.STAGE and x.src[0].arg.addrspace == AddrSpace.LOCAL]
|
||||
exclude_ranges = UOp.group(*[UOp.group(*x.src[1:]) for x in local_indexes]).ranges
|
||||
subs = [(k,v) for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST]
|
||||
# if it's bufferized or a reduce, it's pcontig
|
||||
is_pcontig, is_subs = partition(subs, lambda x: x[0] in exclude_ranges or any([r.arg[-1] == AxisType.REDUCE for r in x[1].ranges]))
|
||||
if not len(is_subs):
|
||||
return None
|
||||
if len(is_pcontig):
|
||||
ret = src.substitute(dict(is_subs), extra_pm=pm_gate_substitute)
|
||||
return ret.bufferize(*[x[0] for x in is_pcontig], arg=BufferizeOpts(None, AddrSpace.LOCAL)).index(*[x[1] for x in is_pcontig])
|
||||
else:
|
||||
return None
|
||||
|
||||
# if it makes it here, the bufferize is removed
|
||||
# this is the ranges replaced
|
||||
# NOTE: if buf src is a const, we don't replace it. if idx is Invalid (dead load), don't replace it either
|
||||
replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.val is Invalid)}
|
||||
return src.substitute(replaced, extra_pm=pm_gate_substitute)
|
||||
|
||||
def remove_noop_bufferize(idx,b2):
|
||||
if idx.src[1:] != b2.src[1:]: return None
|
||||
return idx.src[0].shrink(tuple((0, s) for s in b2.shape)) if b2.shape else idx.src[0]
|
||||
|
||||
def after_all_invalid(after:UOp):
|
||||
buf = after.src[0].buf_uop
|
||||
# check all ranges are used (no expand), and same size (no pad and shrink)
|
||||
return all(s.op is Ops.END and (st:=s.src[0]).op is Ops.STORE and st.src[1].base.is_invalid and st.src[0].buf_uop is buf
|
||||
and all(r in st.src[0].ranges for r in s.ended_ranges)
|
||||
and resolve(cast(UOp, prod(r.src[0] for r in s.ended_ranges)).eq(buf.numel()), False) for s in after.src[1:])
|
||||
|
||||
pm_const_buffer_folding = pm_mops+PatternMatcher([
|
||||
(UPat(Ops.STAGE, name="b"), cleanup_dead_axes),
|
||||
# remove noop buffers. if we look at the next index we can remove even more of these
|
||||
(UPat(Ops.INDEX, name="idx").f(Ops.STAGE, allow_any_len=True, name="b2"), remove_noop_bufferize),
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.STAGE),), allow_any_len=True, name="idx").f(Ops.NOOP).f(Ops.STAGE, allow_any_len=True, name="b2"),
|
||||
remove_noop_bufferize),
|
||||
# no buffers for const (ranges don't matter for const - it's the same value everywhere)
|
||||
(UPat(Ops.CONST, name='c').f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)),
|
||||
# indexing a const is a const
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c),
|
||||
# indexing an after with all fully invalid stores is invalid
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.AFTER, name="after"),), allow_any_len=True, name="idx"),
|
||||
lambda idx,after: idx.const_like(Invalid) if after_all_invalid(after) else None),
|
||||
# hack if a noop turned to a const
|
||||
(UPat(Ops.NOOP, src=(UPat.cvar("c"),)), lambda c: c),
|
||||
# mstack on CONST is CONST
|
||||
(UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True),
|
||||
lambda s: c if (c:=s.base).op is Ops.CONST else None),
|
||||
])
|
||||
|
||||
pm_remove_bufferize = PatternMatcher([
|
||||
# remove reindexing with cost function
|
||||
(UPat.var("src").f(Ops.STAGE, allow_any_len=True, name="buf").f(Ops.INDEX, allow_any_len=True, name="idx"), remove_bufferize),
|
||||
# STORE to self is NOOP
|
||||
(UPat.var("x").store(UPat.var("x")), lambda x: UOp(Ops.NOOP)),
|
||||
# END on NOOP is NOOP
|
||||
(UPat(Ops.END, src=(UPat(Ops.NOOP, name="x"),), allow_any_len=True), lambda x: x),
|
||||
])
|
||||
|
||||
DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8, "CPU": 31} # TODO: get from device?
|
||||
def limit_bufs(ctx:IndexingContext, root:UOp):
|
||||
if (device:=root.device) is None: return None # no device, index related calculations
|
||||
device = device if isinstance(device, str) else device[0].split(":")[0]
|
||||
if not (MAX_BUFS:=MAX_KERNEL_BUFFERS.value or DEVICE_MAX_BUFS.get(device, 0)): return None
|
||||
|
||||
def visitor(u:UOp) -> frozenset[UOp]:
|
||||
if u.op in {Ops.STAGE, Ops.AFTER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK}: return frozenset((u,))
|
||||
if len(u.src) == 1: return ctx.buf_cache[u.src[0]]
|
||||
return frozenset().union(*[ctx.buf_cache[s] for s in u.src])
|
||||
bufs = root.topovisit(visitor, ctx.buf_cache)
|
||||
|
||||
if len(bufs) > MAX_BUFS - 1: # NOTE: this -1 is for the output buffer
|
||||
srcs = []
|
||||
for s in root.src:
|
||||
if s.op in GroupOp.Elementwise and s.device is not None:
|
||||
# Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.WEAK, the DEVICE range stays a launched axis
|
||||
orig_ranges = s.ranges
|
||||
end_ranges = [x.replace(arg=(next(ctx.range_idx), AxisType.WEAK)) if x.op is Ops.RANGE and x.arg[-1] is not AxisType.DEVICE else x
|
||||
for x in s.ranges]
|
||||
s = s.substitute(dict(zip(orig_ranges, end_ranges))).bufferize(*end_ranges, arg=BufferizeOpts(device=s.device)).index(*orig_ranges)
|
||||
srcs.append(s)
|
||||
return root.replace(src=tuple(srcs))
|
||||
pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary), name="root"), limit_bufs)])
|
||||
|
||||
# *****************
|
||||
# 4. put in buffers for bufferize
|
||||
# TODO: should BUFFERIZE look a lot more like STORE
|
||||
# BUFFERIZE has device in arg
|
||||
# BUFFERIZE doesn't have indexing, that's implied by the ranges it closes
|
||||
# BUFFERIZE returns the BUFFER ready for INDEXing (doing this will make splitting a lot easier)
|
||||
# NOTE: this has been fixed up a bit
|
||||
|
||||
def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):
|
||||
size = prod(x.shape)
|
||||
dtype = strong_dtype(x.dtype) # a BUFFER is never weak: store at the concrete dtype, the .cast(x.dtype) on the result keeps readers unchanged
|
||||
rngs = sorted(idx.ranges, key=lambda x: x.arg)
|
||||
assert size > 0 and isinstance(size, int), f"no zero sized or symbolic sized buffers {size}"
|
||||
|
||||
# AFTER: add END to the existing STORE, return buffer with kernel dependency
|
||||
if (after:=x.src[0]).op is Ops.AFTER:
|
||||
buf = after.src[0].buf_uop.base
|
||||
if not (stores := [s for s in after.src[1:] if s.op is Ops.STORE and s.src[0].op is Ops.INDEX]): return buf
|
||||
# BUFFERIZE(INDEX(...)); store through the underlying global index instead.
|
||||
ended_stores = []
|
||||
for store in stores:
|
||||
store_target = store.src[0]
|
||||
if store_target.src[0].op is Ops.STAGE and store_target.src[0].src[0].op is Ops.INDEX:
|
||||
store_target = store_target.src[0].src[0]
|
||||
if store.src[1] is store_target: continue # skip self-assign
|
||||
end_rngs = sorted(dedup(tuple(store_target.ranges) + tuple(rngs)), key=lambda x: x.arg)
|
||||
ended_stores.append(store_target.store(store.src[1]).end(*end_rngs))
|
||||
return buf.after(*ended_stores)
|
||||
|
||||
# NOTE: the local BUFFER needs to be disambiguated here
|
||||
if x.arg.addrspace == AddrSpace.GLOBAL:
|
||||
buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(next(ctx), dtype, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
|
||||
do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs)
|
||||
return buf.after(do_store).cast(x.dtype)
|
||||
|
||||
if allow_locals:
|
||||
# handle locals
|
||||
buf = UOp.placeholder((size,), dtype, next(ctx), AddrSpace.LOCAL)
|
||||
do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs)
|
||||
return buf.after(do_store).cast(x.dtype)
|
||||
|
||||
# collapse any BUFFERIZE to single input BUFFERIZE
|
||||
def flatten_bufferize(x:UOp):
|
||||
if len(x.src) == 2: return None
|
||||
ret = x.replace(src=(x.src[0], get_single_element(apply_movement_op(Ops.RESHAPE, (prod(x.shape),), x.shape, x.src[1:]))))
|
||||
rngs = x.src[1:]
|
||||
ret = ret.reshape(x.shape)
|
||||
if any(r.op is Ops.RANGE and r.src[0].op is not Ops.CONST for r in rngs):
|
||||
sym_shape = tuple([r.src[0] if r.op is not Ops.CONST else 1 for r in rngs])
|
||||
ret = ret.shrink(tuple([(0,x) for x in sym_shape]))
|
||||
return ret
|
||||
pm_flatten_bufferize = PatternMatcher([(UPat(Ops.STAGE, name="x"), flatten_bufferize)])
|
||||
|
||||
def is_noop_after_dep(x:UOp) -> bool:
|
||||
return (x.op is Ops.NOOP and len(x.src) == 0) or (x.op is Ops.END and is_noop_after_dep(x.src[0]))
|
||||
|
||||
def remove_noop_afters(x:UOp) -> UOp|None:
|
||||
src = (x.src[0],) + tuple(s for s in x.src[1:] if not is_noop_after_dep(s))
|
||||
if len(src) != len(x.src): return src[0] if len(src) == 1 else x.replace(src=src)
|
||||
return None
|
||||
|
||||
pm_add_buffers = pm_mops+pm_flatten_bufferize+PatternMatcher([
|
||||
(UPat(Ops.STAGE, src=(UPat(), UPat(name="idx")), name="x"), lambda ctx,x,idx: bufferize_to_store(ctx, x, idx, allow_locals=False)),
|
||||
|
||||
# INDEX of a buffer through the weak cast added above: index the buffer directly and cast the loaded value instead.
|
||||
# this must run in the same rewrite that adds the cast, or the expander expands the whole casted buffer into one big VECTORIZE
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("buf"),)),), allow_any_len=True, name="u"),
|
||||
lambda u,buf: u.replace(dtype=None, src=(buf,)+u.src[1:]).cast(u.dtype)),
|
||||
|
||||
# move RESHAPEs through MSELECT/MSTACK
|
||||
(UPat((Ops.MSELECT, Ops.MSTACK), src=UPat(Ops.RESHAPE), name="m"),
|
||||
lambda m: m.replace(src=tuple([x.src[0].base for x in m.src])).reshape(m.shape)),
|
||||
|
||||
# remove any RESHAPEs on KERNEL
|
||||
(UPat(Ops.CALL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
|
||||
|
||||
# remove invalid writes
|
||||
(UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONTIGUOUS, src=(UPat(Ops.CONST, arg=Invalid),)))), lambda: UOp(Ops.NOOP)),
|
||||
(UPat(Ops.STORE, src=(UPat(), UPat(Ops.CONST, arg=Invalid))), lambda: UOp(Ops.NOOP)),
|
||||
(UPat(Ops.AFTER, name="x"), remove_noop_afters),
|
||||
])
|
||||
|
||||
# *****************
|
||||
# 5. split into kernels
|
||||
|
||||
@dataclass
|
||||
class LocalAddBufferContext:
|
||||
dg:int = 0
|
||||
map:dict = field(default_factory=dict)
|
||||
vars:dict = field(default_factory=dict)
|
||||
range:int = 0
|
||||
opts:tuple|None = None
|
||||
|
||||
def debuf(ctx:LocalAddBufferContext, buf:UOp):
|
||||
param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape)),),
|
||||
arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
|
||||
ret = param.reshape(buf.max_shape)
|
||||
# if the buffer has symbolic shape, shrink the max-sized view to the actual shape
|
||||
if buf.max_shape != buf.shape: ret = ret.shrink(tuple((0, s) for s in buf.shape))
|
||||
if buf not in ctx.map: ctx.map[buf] = buf
|
||||
ctx.dg += 1
|
||||
return ret
|
||||
|
||||
def unbind_kernel(ctx:LocalAddBufferContext, b:UOp):
|
||||
ctx.vars[b] = None
|
||||
return b.src[0]
|
||||
|
||||
def handle_after(ctx:LocalAddBufferContext, after:UOp):
|
||||
if after.addrspace == AddrSpace.LOCAL: return None
|
||||
buf = after.buf_uop
|
||||
# NOTE: this is bottom up, so we only add it once
|
||||
if buf not in ctx.map: ctx.map[buf] = after
|
||||
return buf
|
||||
|
||||
def renumber_range(ctx:LocalAddBufferContext, r:UOp):
|
||||
if r.tag != (): return None
|
||||
ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
|
||||
ctx.range += 1
|
||||
return ret
|
||||
|
||||
def find_bufs(x:UOp):
|
||||
idxs = [s for s in x.toposort(gate=lambda x: x.op is not Ops.AFTER) if s.op is Ops.INDEX]
|
||||
read_from: dict[UOp, Ops] = {}
|
||||
if any((buf:=idx.buf_uop).op in {Ops.BUFFER, Ops.PARAM} and read_from.setdefault(buf, op:=idx.src[0].op) is not op for idx in idxs):
|
||||
raise RuntimeError(f"cycle detected while indexing {buf}")
|
||||
|
||||
to_define_global = PatternMatcher([
|
||||
(UPat(Ops.STORE, name="x"), find_bufs),
|
||||
(UPat((Ops.BUFFER, Ops.MSTACK, Ops.MSELECT), name="buf"), debuf),
|
||||
(UPat(Ops.PARAM, name="v"), lambda v:
|
||||
UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype, multiple_of=v.arg.multiple_of)
|
||||
if v.arg.name is not None and v.arg.vmin_vmax is not None else None),
|
||||
|
||||
# this renumbers the params
|
||||
(UPat(Ops.PARAM, name="buf"), lambda ctx, buf:
|
||||
None if buf.tag != () or buf.arg.name is not None or buf._shape is None else debuf(ctx, buf)),
|
||||
|
||||
# ALU params are scalar symbolic values, not buffers.
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.PARAM, name="v"),)), lambda v: v if v.addrspace == AddrSpace.ALU else None),
|
||||
|
||||
(UPat(Ops.BIND, name="b"), unbind_kernel),
|
||||
(UPat(Ops.AFTER, name="after"), handle_after),
|
||||
|
||||
# remove device from local BUFFERIZE
|
||||
(UPat(Ops.STAGE, name="b"), lambda b: b.replace(arg=replace(b.arg, device=None))),
|
||||
|
||||
# renumber the ranges starting with 0 so that kernel deduping works
|
||||
(UPat(Ops.RANGE, name="r"), renumber_range),
|
||||
])
|
||||
|
||||
def get_contiguous(ctx:LocalAddBufferContext, x:UOp):
|
||||
if isinstance(x.arg, tuple) and all(isinstance(y, Opt) for y in x.arg): ctx.opts = x.arg
|
||||
return x.src[0]
|
||||
|
||||
rangeify_codegen = PatternMatcher([
|
||||
(UPat(Ops.CONTIGUOUS, name="x"), get_contiguous),
|
||||
|
||||
# no NOOP in the kernel graph
|
||||
# TODO: this can be moved into codegen?
|
||||
(UPat(Ops.NOOP, name="x"), lambda x: x.src[0] if len(x.src) else None),
|
||||
])
|
||||
|
||||
pm_add_param_range_tags = PatternMatcher([
|
||||
(UPat((Ops.PARAM, Ops.RANGE), name="x"), lambda x: x.rtag(())),
|
||||
])
|
||||
|
||||
def split_store(x:UOp) -> UOp|None:
|
||||
# if we have any open ranges here, we don't split. open DEVICE ranges are fine, they are bound per device at launch
|
||||
if any(r.arg[-1] is not AxisType.DEVICE for r in x.ranges): return None
|
||||
|
||||
# local kernel rewrite
|
||||
lctx = LocalAddBufferContext()
|
||||
ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)
|
||||
|
||||
# create the Kernel. NOTE: buffers can be on different devices here now, they are compiled to SDMA copies later by schedule
|
||||
return ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)).call(*lctx.map.values(), *lctx.vars.keys())
|
||||
|
||||
split_kernels = PatternMatcher([
|
||||
(UPat((Ops.STORE, Ops.END), name="x"), split_store),
|
||||
])
|
||||
|
||||
def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
|
||||
input_src = copy.src[0]
|
||||
if not input_src.has_buffer_identity(after_ok=True): input_src = input_src.contiguous()
|
||||
input_src = input_src.flatten()
|
||||
if existing_buf is not None:
|
||||
# if the existing buffer is not a full buffer, we can't use it
|
||||
if not existing_buf.has_buffer_identity(after_ok=True): return None
|
||||
# if there's already a buffer, we just use it
|
||||
return existing_buf.flatten().store(input_src)
|
||||
# create the output buffer
|
||||
buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
|
||||
# reshape back to input
|
||||
return buf.after(buf.store(input_src)).reshape(copy.shape)
|
||||
|
||||
pm_copy_to_store = PatternMatcher([
|
||||
(UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
|
||||
(UPat(Ops.COPY, name="copy"), convert_copy_to_store),
|
||||
])
|
||||
|
||||
@rewrite_group(new_ctx=False)
|
||||
def get_kernel_graph(sink:UOp) -> UOp:
|
||||
tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
|
||||
if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
|
||||
tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
|
||||
|
||||
tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
|
||||
|
||||
# convert movement ops to ranges
|
||||
tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
|
||||
|
||||
tsink = graph_rewrite(tsink, symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize,
|
||||
name="symbolic+reduce_collapse+debuf")
|
||||
tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
|
||||
|
||||
if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
|
||||
|
||||
# bufferize -> store
|
||||
slots = [x.arg.slot for x in tsink.toposort() if x.op is Ops.BUFFER and isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL]
|
||||
paramarg_start: int = max([-1]+slots) + 1
|
||||
tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_param_range_tags, ctx=itertools.count(paramarg_start), bottom_up=True, name="stage to store")
|
||||
tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
|
||||
|
||||
if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
|
||||
return tsink
|
||||
Reference in New Issue
Block a user