mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-08 22:48:25 -05:00
add Scheduler to replace Kernel with POSTOPT=2 (#11924)
* ** simple kernel to replace Kernel for postopt * support old * fix beam * beaming * beam on old * bring tensor cores back * raise * postbeam * test ops passes on mac * skip that * postopt default * gate that * fix tensor cores * a few test fixes * dsp fix * tc fix * loop * support swap * test_gemv * fix beam for variable * test opts from high level stuff * range annoying * compile slow * metal slow * better beam * no POSTBEAM * fix nolocals * hc opt mostly works * put that back * lil * some work * fix that * POSTOPT 2 * fix tests * no postopt 2 * work * back * padded tensors cores * shift_to * postopt 0 passes? * write PADTO * fix padded tensor cores * compare hcopt * 18000 lines * should pass tests * fix rangeify * put types back
This commit is contained in:
George Hotz
4
.github/workflows/test.yml
vendored
4
.github/workflows/test.yml
vendored
@@ -382,8 +382,8 @@ jobs:
|
||||
PYTHONPATH=. python extra/optimization/extract_dataset.py
|
||||
gzip -c /tmp/sops > extra/datasets/sops.gz
|
||||
DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
|
||||
- name: Repo line count < 17500 lines
|
||||
run: MAX_LINE_COUNT=17500 python sz.py
|
||||
- name: Repo line count < 18000 lines
|
||||
run: MAX_LINE_COUNT=18000 python sz.py
|
||||
|
||||
fuzzing:
|
||||
name: Fuzzing
|
||||
|
||||
22
extra/test_hcopt.py
Normal file
22
extra/test_hcopt.py
Normal file
@@ -0,0 +1,22 @@
|
||||
from extra.optimization.helpers import load_worlds, ast_str_to_lin
|
||||
from tinygrad.codegen.lowerer import pm_lowerer, get_index
|
||||
from tinygrad.uop.ops import graph_rewrite
|
||||
from tinygrad.codegen.opt.postrange import Scheduler
|
||||
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
|
||||
|
||||
if __name__ == "__main__":
|
||||
ast_strs = load_worlds()
|
||||
for i, ast_str in enumerate(ast_strs):
|
||||
lin = ast_str_to_lin(ast_str)
|
||||
opt1 = hand_coded_optimizations(lin)
|
||||
|
||||
lowered = graph_rewrite(lin.ast, pm_lowerer, ctx=get_index(lin.ast), bottom_up=True)
|
||||
sch = Scheduler(lowered, lin.opts)
|
||||
opt2 = hand_coded_optimizations(sch)
|
||||
|
||||
if opt1 != opt2:
|
||||
print("*******")
|
||||
print("Kernel: ", opt1)
|
||||
print("Scheduler: ", opt2)
|
||||
else:
|
||||
print("******* MATCH")
|
||||
56
test/external/external_metal_compile_slow.py
vendored
Normal file
56
test/external/external_metal_compile_slow.py
vendored
Normal file
@@ -0,0 +1,56 @@
|
||||
# ruff: noqa: E501
|
||||
from tinygrad import dtypes
|
||||
from tinygrad.helpers import Timing, getenv
|
||||
from tinygrad.codegen.opt.kernel import Opt, OptOps
|
||||
from tinygrad.engine.realize import get_program, CompiledRunner
|
||||
from tinygrad.uop.ops import UOp, Ops, AxisType
|
||||
|
||||
if __name__ == "__main__":
|
||||
if getenv("TC", 0) == 0:
|
||||
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1179648), arg=0, src=())
|
||||
c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
|
||||
c2 = UOp.range(UOp.const(dtypes.int, 64), 1, AxisType.GLOBAL)
|
||||
c3 = UOp.range(UOp.const(dtypes.int, 6), 2, AxisType.GLOBAL)
|
||||
c4 = UOp.range(UOp.const(dtypes.int, 6), 3, AxisType.GLOBAL)
|
||||
c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(2097152), arg=1, src=())
|
||||
c6 = UOp.range(UOp.const(dtypes.int, 64), 1004, AxisType.REDUCE)
|
||||
c7 = UOp.range(UOp.const(dtypes.int, 3), 1005, AxisType.REDUCE)
|
||||
c8 = UOp.range(UOp.const(dtypes.int, 3), 1006, AxisType.REDUCE)
|
||||
c9 = c5.index(((((((c1*UOp.const(dtypes.int, 4096))+(c3*UOp.const(dtypes.int, 8)))+c4)+(c6*UOp.const(dtypes.int, 64)))+(c7*UOp.const(dtypes.int, 8)))+c8), UOp.const(dtypes.bool, True)).load()
|
||||
c10 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(36864), arg=2, src=())
|
||||
c11 = c10.index(((((c2*UOp.const(dtypes.int, 576))+(c6*UOp.const(dtypes.int, 9)))+(c7*UOp.const(dtypes.int, 3)))+c8), UOp.const(dtypes.bool, True)).load()
|
||||
c12 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(64), arg=3, src=())
|
||||
c13 = c12.index(c2, UOp.const(dtypes.bool, True)).load()
|
||||
c14 = ((c9*c11).reduce(c6, c7, c8, arg=Ops.ADD)+c13)
|
||||
c15 = c0.index(((((c1*UOp.const(dtypes.int, 2304))+(c2*UOp.const(dtypes.int, 36)))+(c3*UOp.const(dtypes.int, 6)))+c4), UOp.const(dtypes.bool, True)).store(c14, c1, c2, c3, c4)
|
||||
ast = c15.sink()
|
||||
|
||||
# this does have tons of locals
|
||||
opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16), Opt(op=OptOps.UPCAST, axis=3, arg=0),
|
||||
Opt(op=OptOps.LOCAL, axis=0, arg=16), Opt(op=OptOps.UPCAST, axis=3, arg=2),
|
||||
Opt(op=OptOps.GROUPTOP, axis=0, arg=16)]
|
||||
else:
|
||||
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(10616832), arg=0, src=())
|
||||
c1 = UOp.range(UOp.const(dtypes.int, 512), 0, AxisType.GLOBAL)
|
||||
c2 = UOp.range(UOp.const(dtypes.int, 64), 1, AxisType.GLOBAL)
|
||||
c3 = UOp.range(UOp.const(dtypes.int, 36), 2, AxisType.GLOBAL)
|
||||
c4 = UOp.range(UOp.const(dtypes.int, 9), 3, AxisType.GLOBAL)
|
||||
c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(36864), arg=1, src=())
|
||||
c6 = UOp.range(UOp.const(dtypes.int, 64), 1004, AxisType.REDUCE)
|
||||
c7 = c5.index((((c2*UOp.const(dtypes.int, 9))+c4)+(c6*UOp.const(dtypes.int, 576))), UOp.const(dtypes.bool, True)).load()
|
||||
c8 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1179648), arg=2, src=())
|
||||
c9 = c8.index((((c1*UOp.const(dtypes.int, 2304))+c3)+(c6*UOp.const(dtypes.int, 36))), UOp.const(dtypes.bool, True)).load()
|
||||
c10 = (c7*c9).reduce(c6, arg=Ops.ADD)
|
||||
c11 = c0.index(((((c1*UOp.const(dtypes.int, 20736))+(c2*UOp.const(dtypes.int, 324)))+(c3*UOp.const(dtypes.int, 9)))+c4), UOp.const(dtypes.bool, True)).store(c10, c1, c2, c3, c4)
|
||||
ast = c11.sink()
|
||||
|
||||
opts = [Opt(op=OptOps.TC, axis=0, arg=(0, 0, 1)), Opt(op=OptOps.UPCAST, axis=2, arg=4),
|
||||
Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=0)]
|
||||
|
||||
prg = get_program(ast, opts=opts)
|
||||
print(prg.src)
|
||||
for i in range(10):
|
||||
with Timing(f"try {i}: "):
|
||||
# NOTE: this doesn't even run the kernel
|
||||
try: CompiledRunner(prg)
|
||||
except RuntimeError: pass
|
||||
@@ -180,6 +180,7 @@ class TestOuterworld(unittest.TestCase):
|
||||
out.realize()
|
||||
print(out.numpy())
|
||||
|
||||
@unittest.skip("opts don't work")
|
||||
def test_triple_gemm(self):
|
||||
x = Tensor.rand(1, 16).realize()
|
||||
W = Tensor.rand(3, 16, 16).realize()
|
||||
|
||||
@@ -4,6 +4,7 @@ from tinygrad.codegen import full_rewrite_to_sink
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from tinygrad.uop.symbolic import simplify_valid
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
def get_gated_load_uop(valid:UOp, idx:UOp):
|
||||
return UOp(Ops.LOAD, dtypes.float, (
|
||||
@@ -45,7 +46,8 @@ class TestHelpers(unittest.TestCase):
|
||||
|
||||
class TestValidIdxSimplification(unittest.TestCase):
|
||||
def check(self, load, sidx, svalid):
|
||||
load = full_rewrite_to_sink(load.sink()).src[0]
|
||||
with Context(NOOPT=1):
|
||||
load = full_rewrite_to_sink(load.sink()).src[0]
|
||||
idx, valid = load.src[0].src[1], load.src[0].src[2]
|
||||
self.assertEqual(idx.render(simplify=False), sidx)
|
||||
self.assertEqual(valid.render(simplify=False), svalid)
|
||||
@@ -208,7 +210,8 @@ class TestValidIdxSimplification(unittest.TestCase):
|
||||
|
||||
class TestImageSimplification(unittest.TestCase):
|
||||
def check(self, load, svalid, sidx0, sidx1):
|
||||
load = full_rewrite_to_sink(load.sink()).src[0]
|
||||
with Context(NOOPT=1):
|
||||
load = full_rewrite_to_sink(load.sink()).src[0]
|
||||
idx = load.src[0].src[1]
|
||||
self.assertEqual(idx.op, Ops.VECTORIZE)
|
||||
self.assertEqual(len(idx.src), 2)
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
from tinygrad.codegen.opt.kernel import Kernel
|
||||
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, KernelInfo
|
||||
from tinygrad.helpers import NOOPT, BEAM, getenv
|
||||
from tinygrad.helpers import NOOPT, BEAM, getenv, POSTOPT
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.uop.spec import type_verify
|
||||
|
||||
@@ -26,7 +26,7 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp|None:
|
||||
k = Kernel(ast, opts=renderer)
|
||||
if not NOOPT:
|
||||
k.apply_opts(hand_coded_optimizations(k))
|
||||
if BEAM >= 1:
|
||||
if not POSTOPT and BEAM >= 1:
|
||||
from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
|
||||
kb = Kernel(ast, opts=renderer)
|
||||
rawbufs = bufs_from_lin(kb, allocate=False)
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
import itertools
|
||||
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError, AxisType
|
||||
from tinygrad.codegen.opt.postrange import Scheduler
|
||||
from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, AMX
|
||||
from tinygrad.dtype import ImageDType
|
||||
from tinygrad.uop.ops import Ops, resolve
|
||||
|
||||
def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
|
||||
# first try the tensor cores
|
||||
""" Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
|
||||
Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
|
||||
@@ -29,7 +30,7 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
|
||||
|
||||
# skip hand-coded TC opts if AMX, upcasting will make kernel slower
|
||||
if (tc_opts:=tk.tensor_core_opts) is not None and not AMX:
|
||||
if isinstance(k, Kernel) and (tc_opts:=tk.tensor_core_opts) is not None and not AMX:
|
||||
# hand-coded TC opts
|
||||
for tc_dim in [tc_dim for tc_dim in [1,0] if tc_opts.axes_exist[tc_dim]]: # attempt to upcast M and N
|
||||
szs = [sz for sz in [5,4,3,2] if tk.full_shape[tc_opts.axes[tc_dim]] % sz == 0]
|
||||
@@ -49,19 +50,20 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and \
|
||||
k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \
|
||||
(mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
|
||||
st0, st1 = k.sts[k.bufs.index(mulop.src[0])], k.sts[k.bufs.index(mulop.src[1])]
|
||||
strides0, strides1 = st0.real_strides(), st1.real_strides()
|
||||
def has_expanded_axis(shape, strides): return any(resolve(s > 1) and not resolve(st != 0) for s,st in zip(shape,strides))
|
||||
if strides0[first_reduce:=(k.axes_of(AxisType.REDUCE)[0])] == 1 and \
|
||||
not (has_expanded_axis(st0.shape, strides0) and has_expanded_axis(st1.shape, strides1)):
|
||||
for global_idx in k.axes_of(AxisType.GLOBAL):
|
||||
if k.full_shape[first_reduce]%MV_THREADS_PER_ROW == 0 and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
|
||||
if DEBUG >= 3:
|
||||
print(f"MATVEC: {k.full_shape=} {first_reduce=} {strides0=} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
|
||||
if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
|
||||
if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
|
||||
if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
|
||||
return k.applied_opts
|
||||
if isinstance(k, Kernel):
|
||||
st0, st1 = k.sts[k.bufs.index(mulop.src[0])], k.sts[k.bufs.index(mulop.src[1])]
|
||||
strides0, strides1 = st0.real_strides(), st1.real_strides()
|
||||
def has_expanded_axis(shape, strides): return any(resolve(s > 1) and not resolve(st != 0) for s,st in zip(shape,strides))
|
||||
if strides0[first_reduce:=(k.axes_of(AxisType.REDUCE)[0])] == 1 and \
|
||||
not (has_expanded_axis(st0.shape, strides0) and has_expanded_axis(st1.shape, strides1)):
|
||||
for global_idx in k.axes_of(AxisType.GLOBAL):
|
||||
if k.full_shape[first_reduce]%MV_THREADS_PER_ROW == 0 and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
|
||||
if DEBUG >= 3:
|
||||
print(f"MATVEC: {k.full_shape=} {first_reduce=} {strides0=} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
|
||||
if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
|
||||
if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
|
||||
if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
|
||||
return k.applied_opts
|
||||
|
||||
# are we grouping? (requires local shape support)
|
||||
if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
|
||||
@@ -74,7 +76,12 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
# upcast float4 images
|
||||
for buf_index,buf in enumerate(k.bufs):
|
||||
if isinstance(buf.src[0].dtype, ImageDType):
|
||||
if (unit_stride_axes_mul_4 := [i for i in k.sts[buf_index].unit_stride_axes(ignore_valid=True) if k.sts[buf_index].shape[i]%4 == 0]):
|
||||
if hasattr(k, "sts"):
|
||||
unit_stride_axes_mul_4 = [i for i in k.sts[buf_index].unit_stride_axes(ignore_valid=True) if k.sts[buf_index].shape[i]%4 == 0]
|
||||
else:
|
||||
# part of real_strides
|
||||
unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].split_uop(Ops.ADD) if c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
|
||||
if len(unit_stride_axes_mul_4):
|
||||
if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims:
|
||||
k.apply_opt(Opt(OptOps.UPCAST, axis, 4))
|
||||
elif axis in k.unrollable_dims:
|
||||
@@ -89,8 +96,9 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
to_upcast: list[int] = []
|
||||
# upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first)
|
||||
for axis in k.upcastable_dims:
|
||||
if k.full_shape[axis] <= 7 and any(st.axis_is_masked(axis) for st in k.sts) and \
|
||||
prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
|
||||
if isinstance(k, Kernel): is_masked = any(st.axis_is_masked(axis) for st in k.sts)
|
||||
else: is_masked = any(len(st.src) > 2 and k.rngs[axis] in st.src[2].parents for st in k.bufs)
|
||||
if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
|
||||
if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
|
||||
to_upcast.append(axis)
|
||||
for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0))
|
||||
@@ -104,10 +112,24 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
|
||||
# if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
|
||||
if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
|
||||
if any(st.views[-1].strides[axis] == 0 and \
|
||||
all(x != 0 for t,x in zip(k.axis_types, st.real_strides()) if t in (AxisType.UPCAST, AxisType.UNROLL)) for st in k.sts):
|
||||
xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in k.sts),
|
||||
sum(st.views[-1].strides[axis] for st in k.sts), axis, upcast_amount))
|
||||
if isinstance(k, Kernel):
|
||||
# must have stride 0 on a view
|
||||
# must have all non stride 0 on what's upcasted before
|
||||
if any(st.views[-1].strides[axis] == 0 and \
|
||||
all(x != 0 for t,x in zip(k.axis_types, st.real_strides()) if t in (AxisType.UPCAST, AxisType.UNROLL)) for st in k.sts):
|
||||
xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in k.sts),
|
||||
sum(st.views[-1].strides[axis] for st in k.sts), axis, upcast_amount))
|
||||
else:
|
||||
rng = k.rngs[axis]
|
||||
if any(rng not in b.src[1].parents and all(r2 in b.src[1].parents for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs):
|
||||
num_strides, sum_strides = 0, 0
|
||||
for b in k.bufs:
|
||||
if rng in b.src[1].parents: num_strides += 1
|
||||
for c in b.src[1].split_uop(Ops.ADD):
|
||||
if c is rng: sum_strides += 1
|
||||
if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
|
||||
if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
|
||||
xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
|
||||
if xb_choices:
|
||||
xb_choices = sorted(xb_choices)
|
||||
if DEBUG >= 4: print(f"more upcast axis : {xb_choices}")
|
||||
@@ -145,7 +167,11 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
k.apply_opt(Opt(OptOps.NOLOCALS))
|
||||
else:
|
||||
# prioritize making expand axes local
|
||||
local_axis_ranking = [(any(st.views[-1].strides[axis] == 0 for st in k.sts), axis) for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP)]
|
||||
if isinstance(k, Kernel):
|
||||
local_axis_ranking = [(any(st.views[-1].strides[axis] == 0 for st in k.sts), axis) for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP)]
|
||||
else:
|
||||
local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].parents for b in k.bufs), axis) \
|
||||
for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
|
||||
to_local: list[tuple[int, int]] = []
|
||||
for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
|
||||
local_size = prod(sz for _, sz in to_local)
|
||||
|
||||
@@ -1,18 +1,332 @@
|
||||
from dataclasses import replace
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
|
||||
from tinygrad.helpers import colored
|
||||
from tinygrad.codegen.opt.kernel import axis_colors
|
||||
import math, itertools
|
||||
from collections import defaultdict
|
||||
from typing import cast, Final
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, _substitute, AxisType
|
||||
from tinygrad.uop.symbolic import symbolic
|
||||
from tinygrad.device import Buffer
|
||||
from tinygrad.dtype import AddrSpace, dtypes
|
||||
from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up
|
||||
from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.schedule.rangeify import remove_tags
|
||||
|
||||
def rename_sink(s:UOp):
|
||||
if s.arg is not None and s.arg.name != "test": return None
|
||||
# NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
|
||||
axis_to_pos = {AxisType.LOOP: -1, AxisType.GLOBAL: 0, AxisType.LOCAL: 1, AxisType.UPCAST: 2,
|
||||
AxisType.GROUP_REDUCE: 1, AxisType.REDUCE: 3, AxisType.UNROLL: 4}
|
||||
|
||||
# get all ranges (sorted)
|
||||
rngs = sorted([u for u in s.parents if u.op is Ops.RANGE], key=lambda x: x.arg[0:-1])
|
||||
def flatten_range(r:UOp):
|
||||
off = 2 if r.op is Ops.STORE else 1
|
||||
rngs = r.src[off:]
|
||||
if not len(rngs): return None
|
||||
new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
|
||||
return r.replace(src=r.src[:off]+tuple(new_rngs))
|
||||
|
||||
# add name to kernel
|
||||
name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[-1]]) for x in rngs])
|
||||
return s.replace(arg=KernelInfo(name=name) if s.arg is None else replace(s.arg, name=name))
|
||||
pm_flatten_range = PatternMatcher([
|
||||
# real ranges only
|
||||
(UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range),
|
||||
])
|
||||
|
||||
def count_divmod(x:UOp): return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
|
||||
|
||||
class Scheduler:
|
||||
def __init__(self, ast:UOp, opts:Renderer):
|
||||
self.ast, self.opts = ast, opts
|
||||
self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
|
||||
self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
|
||||
|
||||
@property
|
||||
def rngs(self):
|
||||
# always in order by axistype
|
||||
return sorted([u for u in self.ast.parents if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
|
||||
@property
|
||||
def shape_len(self): return len(self.rngs)
|
||||
@property
|
||||
def full_shape(self): return [x.vmax+1 for x in self.rngs]
|
||||
@property
|
||||
def axis_types(self): return [x.arg[-1] for x in self.rngs]
|
||||
@property
|
||||
def maxarg(self): return max([x.arg[0] for x in self.rngs], default=0)
|
||||
|
||||
# strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
|
||||
def shape_str(self) -> list[str]:
|
||||
ret: list[str] = []
|
||||
cnt: dict[AxisType, int] = {}
|
||||
for x in self.axis_types:
|
||||
cnt[x] = (cnt[x] + 1) if x in cnt else 0
|
||||
ret.append(f"{axis_letters[x]}{cnt[x]}")
|
||||
return ret
|
||||
def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])
|
||||
|
||||
@property
|
||||
def termination(self):
|
||||
terminators = [u for u in self.ast.parents if u.op in {Ops.REDUCE, Ops.STORE}]
|
||||
termination = {}
|
||||
for t in terminators:
|
||||
# works without pm_flatten_range
|
||||
for u in UOp.sink(*t.src[1 if t.op is Ops.REDUCE else 2:]).parents:
|
||||
if u.op is Ops.RANGE: termination[u] = t
|
||||
return termination
|
||||
|
||||
def copy(self): return Scheduler(self.get_optimized_ast(), self.opts)
|
||||
|
||||
kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
|
||||
def get_optimized_ast(self, name_override:str|None=None):
|
||||
if name_override is not None: name = name_override
|
||||
else:
|
||||
name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
|
||||
Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
|
||||
num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
|
||||
name += colored(num, 'BLACK')
|
||||
self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
|
||||
return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)
|
||||
|
||||
def convert_loop_to_global(self):
|
||||
if not self.opts.has_local: return None
|
||||
store_rngs = self.ast.src[0].src[2:]
|
||||
|
||||
# filter any not in local stores
|
||||
local_store_rngs = [x.ranges for x in self.ast.toposort() if (x.op is Ops.STORE and x.src[0].ptrdtype.addrspace == AddrSpace.LOCAL) \
|
||||
or (x.op is Ops.BUFFERIZE and x.arg == AddrSpace.LOCAL)]
|
||||
for ls in local_store_rngs: store_rngs = tuple([x for x in store_rngs if x in ls])
|
||||
|
||||
store_rng = [x for x in UOp.sink(*store_rngs).toposort() if x.op is Ops.RANGE] if store_rngs else []
|
||||
rng = [x.replace(arg=(x.arg[0], AxisType.GLOBAL)) if x.arg[1] == AxisType.LOOP and x in store_rng else x for x in self.rngs]
|
||||
|
||||
self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))
|
||||
|
||||
def simplify_merge_adjacent(self):
|
||||
i = 0
|
||||
while i < len(self.rngs)-1:
|
||||
r0, r1 = self.rngs[i], self.rngs[i+1]
|
||||
# same axistype and same termination
|
||||
termination = self.termination
|
||||
if r0.arg[1] == r1.arg[1] and r0 in termination and r1 in termination and termination[r0] == termination[r1]:
|
||||
s0, s1 = r0.src[0], r1.src[0]
|
||||
new_range = r0.replace(src=(s0*s1,)).simplify()
|
||||
# this checks the legality of a merge
|
||||
oidx = self.ast.simplify()
|
||||
nidx = graph_rewrite(oidx, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1}, name=f"check_merge_{i}_{i+1}")
|
||||
# it simplifies
|
||||
if count_divmod(nidx) <= count_divmod(oidx):
|
||||
# it is correct
|
||||
midx = graph_rewrite(nidx, _substitute+symbolic+pm_flatten_range, ctx={new_range:r0*s1+r1}, name=f"correct_merge_{i}_{i+1}")
|
||||
if oidx is midx:
|
||||
self.ast = nidx
|
||||
continue
|
||||
i += 1
|
||||
|
||||
def colors(self) -> list[str]: return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types]
|
||||
def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():4s}', color) for x,color in zip(self.rngs, self.colors())])
|
||||
|
||||
def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False):
|
||||
if (old_sz:=rng.src[0].divides(amount)) is None:
|
||||
raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
|
||||
new_rng = UOp.range(amount, self.maxarg+1, new_type)
|
||||
replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
|
||||
sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
|
||||
self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[0]} {amount}")
|
||||
return replaced_rng, new_rng
|
||||
|
||||
def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
|
||||
def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]
|
||||
@property
|
||||
def upcastable_dims(self): return self.axes_of(AxisType.GLOBAL, AxisType.LOCAL)
|
||||
@property
|
||||
def unrollable_dims(self): return self.axes_of(AxisType.REDUCE, AxisType.GROUP_REDUCE)
|
||||
|
||||
def real_axis(self, op:OptOps, axis:int|None):
|
||||
try:
|
||||
if axis is None: return -1
|
||||
if op is OptOps.UNROLL: return self.unrollable_dims[axis]
|
||||
if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
|
||||
check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
|
||||
return axis
|
||||
except IndexError as e: raise KernelOptError from e
|
||||
|
||||
def apply_opt(self, opt:Opt, append_opt:bool=True):
|
||||
if opt.op is OptOps.NOLOCALS:
|
||||
check(all(x not in {AxisType.LOCAL, AxisType.GROUP_REDUCE} for x in self.axis_types), "no locals can't have locals")
|
||||
self.dont_use_locals = True
|
||||
self.applied_opts.append(opt)
|
||||
return
|
||||
|
||||
if opt.op in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}:
|
||||
check(self.opts.has_local, "locals needed for opt")
|
||||
|
||||
rng = self.rngs[self.real_axis(opt.op, opt.axis)]
|
||||
|
||||
opt_to_at = {
|
||||
OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST,
|
||||
OptOps.UNROLL: AxisType.UNROLL, OptOps.GROUP: AxisType.GROUP_REDUCE,
|
||||
OptOps.GROUPTOP: AxisType.GROUP_REDUCE}
|
||||
|
||||
if opt.op in opt_to_at:
|
||||
amt:int = (rng.vmax+1) if opt.arg == 0 else cast(int, opt.arg)
|
||||
if opt.op is OptOps.UNROLL:
|
||||
check(amt <= 32, "don't unroll more than 32")
|
||||
check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
|
||||
if opt.op is OptOps.UPCAST:
|
||||
check((self.opts is not None and self.opts.device == "DSP") or amt <= 16, "don't upcast more than 16")
|
||||
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, "upcast is for GLOBAL/LOCAL/LOOP")
|
||||
if opt.op is OptOps.LOCAL:
|
||||
check(not self.dont_use_locals, "can't use locals")
|
||||
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
|
||||
if opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:
|
||||
check(not self.dont_use_locals, "can't use locals")
|
||||
check(rng.arg[-1] == AxisType.REDUCE, "group is for reduce")
|
||||
self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op==OptOps.GROUPTOP)
|
||||
elif opt.op is OptOps.TC:
|
||||
check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
|
||||
check(opt.axis is not None, "tensor core opts must have an axis")
|
||||
check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg")
|
||||
check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.opts.tensor_cores), "tensor core opts must have valid tc_select")
|
||||
check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt")
|
||||
check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
|
||||
check(self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt), "no tensor core available")
|
||||
elif opt.op is OptOps.PADTO:
|
||||
check(rng.src[0].op is Ops.CONST, "only pad const")
|
||||
replaced_rng = UOp.range(round_up(rng.vmax+1, cast(int, opt.arg)), *rng.arg)
|
||||
replaces = {rng:replaced_rng}
|
||||
for b in self.bufs:
|
||||
if rng in b.src[1].sparents:
|
||||
valid = replaced_rng < rng.vmax+1
|
||||
if len(b.src) > 2: valid = b.src[2] & valid
|
||||
replaces[b] = b.replace(src=b.src[0:2]+(valid,))
|
||||
self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
|
||||
elif opt.op is OptOps.SWAP:
|
||||
try:
|
||||
altrng = self.rngs[opt.arg]
|
||||
except IndexError:
|
||||
raise KernelOptError
|
||||
check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "swap only for globals")
|
||||
self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1]), tag=1),
|
||||
altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)})
|
||||
self.ast = graph_rewrite(self.ast, remove_tags)
|
||||
else:
|
||||
raise KernelOptError(f"unsupported opt {opt.op}")
|
||||
if append_opt:
|
||||
self.applied_opts.append(opt)
|
||||
|
||||
def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> bool:
|
||||
reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
|
||||
if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore")
|
||||
reduceop = reduceops[0]
|
||||
if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD:
|
||||
mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
|
||||
if mul.op is not Ops.MUL: return False
|
||||
in0, in1 = mul.src
|
||||
try:
|
||||
tensor_cores = self.opts.tensor_cores if tc_select == -1 else [self.opts.tensor_cores[tc_select]]
|
||||
except IndexError:
|
||||
raise KernelOptError(f"invalid tensor core choice {tc_select}")
|
||||
for tc in tensor_cores:
|
||||
if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
|
||||
# tensor cores have three ranges. X, Y, and REDUCE
|
||||
in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0])
|
||||
in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0])
|
||||
red_ranges = sorted(reduceop.src[1:], key=lambda x: x.arg[0])
|
||||
if DEBUG >= 3:
|
||||
print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
|
||||
f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")
|
||||
if not len(in0_ranges) or not len(in1_ranges) or not len(red_ranges): continue
|
||||
|
||||
# pick ranges
|
||||
# NOTE: why are in1 and in0 switched?
|
||||
axis_choices = list(itertools.product(in1_ranges, in0_ranges, red_ranges))
|
||||
if not (axis < len(axis_choices)): continue
|
||||
axes = list(axis_choices[axis])
|
||||
|
||||
# do optimizations and save the ranges
|
||||
try:
|
||||
for i,a in enumerate(axes):
|
||||
# apply_opt should return the updated range?
|
||||
idx = self.rngs.index(a)
|
||||
self.apply_opt(Opt(OptOps.PADTO, idx, tc.dims[i]), append_opt=False) # PADTO might fail
|
||||
axes[i] = self.rngs[idx]
|
||||
except KernelOptError: continue
|
||||
|
||||
ne: list[UOp] = []
|
||||
for opt in tc.opts:
|
||||
axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, {"u":AxisType.UPCAST, "l":AxisType.LOCAL}[opt[0]])
|
||||
ne.append(new_range)
|
||||
for _, amt in tc.get_reduce_axes():
|
||||
axes[2], new_range = self.shift_to(axes[2], amt, AxisType.UNROLL)
|
||||
ne.append(new_range)
|
||||
|
||||
if use_tensor_cores != 2:
|
||||
# fix the srcs
|
||||
reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0]
|
||||
tne = [x.replace(tag=1) for x in ne]
|
||||
ret = reduceop.substitute(dict(zip(ne, tne)))
|
||||
srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
|
||||
srcs = [x.substitute(dict(zip(tne, [ne[i] for i in argsort(p)]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
|
||||
|
||||
# get reduce/upcast axes for the tensor cores
|
||||
tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
|
||||
base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
|
||||
tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
|
||||
|
||||
# axes to range number (was done in lowerer)
|
||||
tc_upcast_axes = tuple([tuple([(self.rngs[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
|
||||
tc_reduce_axes = tuple([self.rngs[a].arg[0] for a in tc_reduce_axes])
|
||||
|
||||
# construct the op
|
||||
# TODO: remove tc_upcast_axes from the arg
|
||||
# do the reduce_axes always disappear? i think they don't
|
||||
# they need to be moved into the WMMA srcs
|
||||
wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
|
||||
wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
|
||||
UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0], tag=1),
|
||||
UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1], tag=1),
|
||||
UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1)
|
||||
tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2], tag=1)
|
||||
|
||||
# preserve extra reduces
|
||||
reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
|
||||
if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), Ops.ADD)
|
||||
self.ast = self.ast.substitute({reduceop: tc_uop})
|
||||
return True
|
||||
return False
|
||||
|
||||
# helpers for hand_coded_optimizations
|
||||
@property
|
||||
def reduceop(self) -> UOp|None:
|
||||
red = [x for x in self.ast.parents if x.op is Ops.REDUCE]
|
||||
if not len(red): return None
|
||||
return UOp(Ops.REDUCE_AXIS, red[0].dtype, red[0].src, (red[0].arg, ()))
|
||||
@property
|
||||
def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]
|
||||
@property
|
||||
def output_shape(self):
|
||||
return [s if at not in {AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE} else 1 for s,at in zip(self.full_shape, self.axis_types)]
|
||||
@property
|
||||
def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
|
||||
@property
|
||||
def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
|
||||
|
||||
def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
|
||||
glbls = sorted([x for x in ast.parents if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
|
||||
return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls]
|
||||
|
||||
def apply_opts(ctx:Renderer, ast:UOp):
|
||||
if ast.tag is not None: return None
|
||||
k = Scheduler(ast, ctx)
|
||||
k.convert_loop_to_global()
|
||||
if BEAM >= 1:
|
||||
k.simplify_merge_adjacent()
|
||||
from tinygrad.codegen.opt.search import beam_search
|
||||
rawbufs = bufs_from_ast(ast, ctx.device)
|
||||
k = beam_search(k, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
|
||||
elif ast.arg is not None and ast.arg.opts_to_apply is not None:
|
||||
for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
|
||||
elif not NOOPT:
|
||||
k.simplify_merge_adjacent()
|
||||
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
|
||||
# NOTE: hand_coded_optimizations doesn't support multiblock opts yet
|
||||
if all(len(u.src) == 1 for u in ast.parents if u.op is Ops.LOAD):
|
||||
for opt in hand_coded_optimizations(k): k.apply_opt(opt)
|
||||
return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)
|
||||
|
||||
pm_postrange_opt = PatternMatcher([
|
||||
(UPat(Ops.SINK, name="s"), rename_sink),
|
||||
(UPat(Ops.SINK, name="ast"), apply_opts),
|
||||
])
|
||||
|
||||
@@ -8,6 +8,7 @@ from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, di
|
||||
from tinygrad.helpers import IGNORE_BEAM_CACHE
|
||||
from tinygrad.dtype import ImageDType, PtrDType
|
||||
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError
|
||||
from tinygrad.codegen.opt.postrange import Scheduler
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.engine.realize import CompiledRunner, get_program
|
||||
from tinygrad.renderer import ProgramSpec
|
||||
@@ -93,6 +94,7 @@ def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_
|
||||
# *** external API ***
|
||||
|
||||
# get (scrap) buffers for timing the linearizer
|
||||
# NOTE: there's also bufs_from_ast in postrange
|
||||
def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]:
|
||||
bufsts: defaultdict[int, list[UOp]] = defaultdict(list)
|
||||
for x in lin.bufs:
|
||||
@@ -110,7 +112,7 @@ def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]:
|
||||
return cast(list[Buffer], rawbufs)
|
||||
|
||||
# get dictionary of all possible actions
|
||||
def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Kernel]:
|
||||
def get_kernel_actions(lin:Kernel|Scheduler, include_0=True, candidates:list[Opt]|None=None) -> dict[int, Kernel|Scheduler]:
|
||||
acted_lins, max_up, max_lcl = {0:lin} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256), getenv("BEAM_LOCAL_MAX", 1024)
|
||||
kernel_actions = (actions if candidates is None else candidates).copy()
|
||||
|
||||
@@ -122,7 +124,7 @@ def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=Non
|
||||
lin2 = lin.copy()
|
||||
try:
|
||||
lin2.apply_opt(a)
|
||||
up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if (tc:=lin2.tensor_core) else 1
|
||||
up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if hasattr(lin2, 'tensor_core') and (tc:=lin2.tensor_core) else 1
|
||||
for s,c in zip(lin2.full_shape, lin2.axis_types):
|
||||
if c in (AxisType.UPCAST, AxisType.UNROLL): up *= s
|
||||
elif c in (AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= s
|
||||
@@ -134,7 +136,7 @@ def get_kernel_actions(lin:Kernel, include_0=True, candidates:list[Opt]|None=Non
|
||||
return acted_lins
|
||||
|
||||
beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
|
||||
def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value) -> Kernel:
|
||||
def beam_search(lin:Kernel|Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
|
||||
global beam_pool
|
||||
key = {"ast": lin.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": lin.opts.device, "suffix": lin.opts.suffix}
|
||||
if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
|
||||
@@ -142,7 +144,7 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True,
|
||||
for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
|
||||
return ret
|
||||
|
||||
beam: list[tuple[Kernel, float]] = [(lin, float("inf"))]
|
||||
beam: list[tuple[Kernel|Scheduler, float]] = [(lin, float("inf"))]
|
||||
seen_libs = set()
|
||||
|
||||
default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
|
||||
@@ -163,8 +165,8 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True,
|
||||
exiting, st = False, time.perf_counter()
|
||||
dev = Device[lin.opts.device]
|
||||
while not exiting:
|
||||
acted_lins: list[Kernel] = flatten([get_kernel_actions(lin, include_0=False).values() for lin,_ in beam])
|
||||
timed_lins: list[tuple[Kernel, float]] = []
|
||||
acted_lins: list[Kernel|Scheduler] = flatten([get_kernel_actions(lin, include_0=False).values() for lin,_ in beam])
|
||||
timed_lins: list[tuple[Kernel|Scheduler, float]] = []
|
||||
_compile_fn = functools.partial(_try_compile_linearized_w_idx, compiler=dev.compiler)
|
||||
least_compute_ops = math.inf
|
||||
for i,proc in (map(_compile_fn, enumerate(acted_lins)) if beam_pool is None else beam_pool.imap_unordered(_compile_fn, enumerate(acted_lins))):
|
||||
|
||||
Reference in New Issue
Block a user