sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

remove opts arg in get_program, everything uses opts_to_apply [pr] (#15767)

Browse Source 
* check Ops.BEAM in process replay

* remove opts from the get_program api

* lint

* simplify

* cleanup
This commit is contained in:
qazal
2026-04-16 22:42:43 +03:00
committed by GitHub
parent f0c12a2004
commit 12c653a743
12 changed files with 92 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/gemm/tinygrad_nv_matmul.py
+1 -1
View File
@@ -29,7 +29,7 @@ if __name__ == "__main__":
Opt(op=OptOps.LOCAL, axis=0, amt=2),
]
k.apply_opts(opts)
prg = get_program(k.ast, k.opts, k.applied_opts)
prg = get_program(k.ast.replace(arg=replace(k.ast.arg, opts_to_apply=tuple(k.applied_opts))), k.opts)
new_src = prg.src
# can mod source here
prg = replace(prg, src=new_src)
test/backend/test_linearizer.py
+25 -22
View File
@@ -11,6 +11,7 @@ from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, getenv
from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
from tinygrad.renderer.ptx import PTXRenderer
from tinygrad.renderer.cstyle import CUDARenderer
from test.helpers import replace_opts
MOCKGPU = getenv("MOCKGPU")
from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
@@ -44,7 +45,7 @@ class TestLinearizer(unittest.TestCase):
tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
ast = helper_linearizer_opt(out)
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
@unittest.expectedFailure
@@ -52,7 +53,7 @@ class TestLinearizer(unittest.TestCase):
tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
ast = helper_linearizer_opt(out)
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -62,7 +63,7 @@ class TestLinearizer(unittest.TestCase):
b = Tensor.empty(16)
out = img.conv2d(w, b)
ast = helper_linearizer_opt(out)
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
# slice at the last loop end
uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
# only valid test if outermost range is the reduce
@@ -83,7 +84,7 @@ class TestLinearizer(unittest.TestCase):
a = Tensor.randn(2, ).realize()
out = a.reshape(2, 1).expand(2, 3).sum()
ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
@@ -91,7 +92,7 @@ class TestLinearizer(unittest.TestCase):
a = Tensor.randn(2, ).realize()
out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
@@ -99,7 +100,7 @@ class TestLinearizer(unittest.TestCase):
a = Tensor([2, 2]).realize()
out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
ast = helper_linearizer_opt(out, wanna_output=[24])
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
# RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
@@ -112,7 +113,7 @@ class TestLinearizer(unittest.TestCase):
b = Tensor.randn(1, 1).realize()
out = (a + b[0]).sum() + b[0]
ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
# LOAD -> RANGE -> LOAD -> STORE
assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -122,7 +123,7 @@ class TestLinearizer(unittest.TestCase):
b = Tensor.randn(1, 1).realize()
out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
assert len(ranges) == 1 # NOTE: it collapses now
@@ -133,7 +134,7 @@ class TestLinearizer(unittest.TestCase):
# these are of size 3 to avoid float4 coalesce
r = a[:-1] + a[1:]
uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
assert num_loads <= 4, "more load uops than needed"
assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -145,7 +146,7 @@ class TestLinearizer(unittest.TestCase):
a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
r = a.expand([2]) + b.expand([2])
uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
assert num_ops <= 1, "more alu uops than needed"
@@ -154,8 +155,8 @@ class TestLinearizer(unittest.TestCase):
x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
r = Tensor.conv2d(x,w,padding=1).relu()
uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer,
opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).uops
accs = [u for u in uops if u.op is Ops.DEFINE_REG]
stores = [u for u in uops if u.op is Ops.STORE]
assert len(accs) == 0 # it's removed now
@@ -167,7 +168,7 @@ class TestLinearizer(unittest.TestCase):
@unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
def test_upcast_with_locals_cpu(self):
out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
prg = get_program(replace_opts(out.schedule()[-1].ast, [Opt(OptOps.LOCAL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
self.assertEqual(len(prg.src.split("for")), 5)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -178,7 +179,7 @@ class TestLinearizer(unittest.TestCase):
x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
r = (x@y).relu()
opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
program = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
program = get_program(replace_opts(r.schedule()[-1].ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
@@ -192,7 +193,7 @@ class TestLinearizer(unittest.TestCase):
def test_zero_fold(self):
a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
r = Tensor.stack(a, b)
uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
assert num_ops == 0, "more alu uops than needed"
@@ -202,14 +203,14 @@ class TestLinearizer(unittest.TestCase):
if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
realized_ast = a.schedule()[-1].ast
program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
assert local[0].dtype.base == acc_dtype
def test_arg_acc_dtype(self):
def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
realized_ast = c.schedule()[-1].ast
program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
self.assertEqual(local[0].dtype.base, expected_dtype)
@@ -237,7 +238,7 @@ class TestLinearizer(unittest.TestCase):
opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
ast = helper_linearizer_opt(r, [opt])
# the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -257,7 +258,7 @@ class TestLinearizer(unittest.TestCase):
# shrink so that the dims do not collapse
t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
ast = helper_linearizer_opt(t+1)
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
idxs = sorted(idxs, key=lambda uop: uop.arg)
assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
@@ -290,7 +291,7 @@ class TestLinearizer(unittest.TestCase):
sched_copy = sched[:]
run_schedule(sched)
np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
program = get_program(sched_copy[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=())
program = get_program(replace_opts(sched_copy[-1].ast, []), renderer=Device[Device.DEFAULT].renderer)
assert not any(u.op == Ops.WHERE for u in program.uops), "found where where where should be folded"
def test_phi_simplification(self):
@@ -352,7 +353,7 @@ class TestLinearizer(unittest.TestCase):
Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
ast = helper_linearizer_opt(out, opts=[opt])
def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
barrier = [u for u in uops if u.op is Ops.BARRIER]
@@ -420,7 +421,9 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
device = real_bufs[0].device
wanna_output = [np.array(x).flatten() for x in wanna_output]
def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
def get_prg(opts):
ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
return CompiledRunner(replace(get_program(ast, renderer=Device[Device.DEFAULT].renderer), device=device))
def check_opt(opts):
prg = get_prg(opts=opts)
test/backend/test_opt_gemm.py
+2 -1
View File
@@ -5,6 +5,7 @@ from tinygrad.helpers import get_single_element
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.schedule import ExecItem
from test.helpers import replace_opts
class TestOptGemm(unittest.TestCase):
@classmethod
@@ -19,7 +20,7 @@ class TestOptGemm(unittest.TestCase):
t = self.a.T @ self.b.T
# TODO: this should be a generic test helper
si = get_single_element(t.schedule())
run = CompiledRunner(get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts))
run = CompiledRunner(get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer))
ExecItem(si.ast, list(si.bufs), prg=run).run()
test = si.bufs[0].numpy().reshape(self.res.shape)
np.testing.assert_allclose(self.res, test, atol=1e-4)
test/backend/test_quantize_onnx.py
+2 -1
View File
@@ -7,6 +7,7 @@ from tinygrad.uop.ops import Ops
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.schedule import ExecItem
from test.helpers import replace_opts
N = 512
@@ -39,7 +40,7 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=
def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
si = out.schedule()[-1]
prg = get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
prg = get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer)
if replace_src is not None:
old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))
test/device/test_hcq.py
+2 -2
View File
@@ -1,7 +1,7 @@
import unittest, ctypes, struct, os, random, numpy as np, time
from tinygrad import Device, Tensor, dtypes
from tinygrad.helpers import getenv, mv_address, DEBUG, DEV
from test.helpers import slow
from test.helpers import slow, replace_opts
from tinygrad.device import Buffer, BufferSpec
from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
from tinygrad.runtime.autogen import libc
@@ -165,7 +165,7 @@ class TestHCQ(unittest.TestCase):
b = a + 1
si = b.schedule()[-1]
runner = CompiledRunner(get_program(si.ast, TestHCQ.d0.renderer, opts=[Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]))
runner = CompiledRunner(get_program(replace_opts(si.ast, [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer))
zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
test/external/process_replay/process_replay.py
Vendored
+9 -10
View File
@@ -10,9 +10,8 @@ if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
try:
from tinygrad.renderer import Renderer, ProgramSpec
from tinygrad.engine.realize import get_program
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.codegen.opt import Opt
from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, BEAM
from tinygrad.uop.ops import UOp, Ops
from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
except ImportError as e:
print(repr(e))
exit(int(ASSERT_DIFF))
@@ -42,13 +41,13 @@ class ProcessReplayWarning(Warning): pass
# *** replay the function and convert return values to string
def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
if ast.op is Ops.BEAM: ast = ast.src[0]
# the ast.arg is non None if we are inside of search.py
sink_arg = ast.arg or KernelInfo()
if opts is not None: sink_arg = replace(sink_arg, opts_to_apply=tuple(opts))
elif BEAM >= 1 and sink_arg.opts_to_apply is None: sink_arg = replace(sink_arg, opts_to_apply=p.applied_opts)
input_ast = ast if ast.op is Ops.PROGRAM else ast.replace(arg=replace(sink_arg, name=p.name))
def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
if ast.op is Ops.PROGRAM: input_ast = ast
else:
sink = ast.src[0] if ast.op is Ops.BEAM else ast
sink_arg = sink.arg
if ast.op is Ops.BEAM: sink_arg = replace(sink_arg, opts_to_apply=p.applied_opts)
input_ast = sink.replace(arg=replace(sink_arg, name=p.name))
p2 = get_program(input_ast, renderer=renderer)
def to_str(ret:ProgramSpec) -> str:
# PYTHON renderer pickles UOps, first unpickle and decode here
test/helpers.py
+3
View File
@@ -1,4 +1,5 @@
import os, time, struct, functools, unittest
from dataclasses import replace
from typing import Any, Callable
import numpy as np
from tinygrad import Tensor, dtypes, Device
@@ -23,6 +24,8 @@ def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
return line_rewrite(linearize(full_sink), pm_linearize_cleanups)
def replace_opts(ast:UOp, opts:list) -> UOp: return ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))
def derandomize_model(model):
for p in get_parameters(model):
p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
test/null/test_process_replay.py
+10 -7
View File
@@ -3,6 +3,7 @@ from tinygrad import Tensor, Device, Context
from tinygrad.engine.realize import get_program
from tinygrad.codegen.opt import Opt, OptOps
from test.external.process_replay.process_replay import replay_get_program
from test.helpers import replace_opts
N = 16
class TestProcessReplay(unittest.TestCase):
@@ -19,22 +20,24 @@ class TestProcessReplay(unittest.TestCase):
def test_replay_empty_opts(self):
# opts=[] means explicitly apply zero opts (unoptimized)
p = get_program(self.ast, self.renderer, opts=[])
good, compare, _ = replay_get_program(p, self.ast, self.renderer, opts=[])
ast = replace_opts(self.ast, [])
p = get_program(ast, self.renderer)
good, compare, _ = replay_get_program(p, ast, self.renderer)
self.assertEqual(good, compare)
def test_replay_with_opt(self):
# opts=[Opt(...)] means apply a specific opt
opts = [Opt(OptOps.UPCAST, 0, 4)]
p = get_program(self.ast, self.renderer, opts=opts)
good, compare, _ = replay_get_program(p, self.ast, self.renderer, opts=opts)
ast = replace_opts(self.ast, opts)
p = get_program(ast, self.renderer)
good, compare, _ = replay_get_program(p, ast, self.renderer)
self.assertEqual(good, compare)
@Context(BEAM=1)
def test_beam(self):
si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1]
with Context(BEAM=1):
si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1]
p = get_program(si.ast, self.renderer)
good, compare, _ = replay_get_program(p, self.ast, self.renderer)
good, compare, _ = replay_get_program(p, si.ast, self.renderer)
self.assertEqual(good, compare)
if __name__ == '__main__':
test/null/test_uops_stats.py
+12 -10
View File
@@ -9,6 +9,7 @@ from tinygrad.dtype import dtypes
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
from tinygrad.device import Device
from tinygrad.renderer.ptx import PTXRenderer
from test.helpers import replace_opts
def flops_mem(uops, ignore_indexing=False):
est = Estimates.from_uops(uops, ignore_indexing)
@@ -175,13 +176,14 @@ class TestStatsOptimized(unittest.TestCase):
self.assertEqual(p.estimates.mem, 3*N*N*4) # 3 NxN mats with floats
def test_gemm(self):
p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[])
p = get_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
self.check_gemm(p)
self.assertEqual(p.estimates.lds, 2*N*N*N*4 + 4*N*N)
def test_gemm_tc_unroll(self):
try:
p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)])
p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
renderer=Device[Device.DEFAULT].renderer)
except KernelOptError:
raise unittest.SkipTest("no tensor cores")
print(p.src)
@@ -190,20 +192,20 @@ class TestStatsOptimized(unittest.TestCase):
# this is a good lesson about why UPCASTing is a good idea
def test_gemm_one_upcasted(self):
p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4)])
p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
self.check_gemm(p)
self.assertEqual(p.estimates.lds, N*N*N*4 + N*N*N*4//4 + 4*N*N)
def test_gemm_upcasted(self):
p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer,
opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)])
p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.UNROLL, 0, 4)]),
renderer=Device[Device.DEFAULT].renderer)
self.check_gemm(p)
self.assertEqual(p.estimates.lds, 2*N*N*N*4//4 + 4*N*N)
def test_gemm_upcasted_locals(self):
try:
p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4),
Opt(OptOps.LOCAL, 0, 4), Opt(OptOps.LOCAL, 1, 4)])
p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 0, 4),
Opt(OptOps.LOCAL, 1, 4)]), renderer=Device[Device.DEFAULT].renderer)
except KernelOptError:
raise unittest.SkipTest("no locals")
self.check_gemm(p)
@@ -211,7 +213,7 @@ class TestStatsOptimized(unittest.TestCase):
def test_gemm_group(self):
try:
p = get_program(self.ast_gemm, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 4)])
p = get_program(replace_opts(self.ast_gemm, [Opt(OptOps.GROUP, 0, 4)]), renderer=Device[Device.DEFAULT].renderer)
except KernelOptError:
raise unittest.SkipTest("no locals")
SZ = N*N*4
@@ -220,14 +222,14 @@ class TestStatsOptimized(unittest.TestCase):
self.assertEqual(p.estimates.lds, 2*N*N*N*4 + SZ*4 + (SZ*4 + 4*N*N)*4)
def test_reduce(self):
p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[])
p = get_program(replace_opts(self.ast_reduce, []), renderer=Device[Device.DEFAULT].renderer)
print(p.name, p.estimates.ops, p.estimates.mem, p.estimates.lds)
self.assertEqual(p.estimates.ops, N*N)
self.assertEqual(p.estimates.mem, N*N*4 + 4)
def test_reduce_group(self):
try:
p = get_program(self.ast_reduce, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.GROUP, 0, 50)])
p = get_program(replace_opts(self.ast_reduce, [Opt(OptOps.GROUP, 0, 50)]), renderer=Device[Device.DEFAULT].renderer)
except KernelOptError:
raise unittest.SkipTest("no locals")
# NOTE: these are wrong, they don't respect the if statement
test/opt/test_gen_float4.py
+16 -15
View File
@@ -4,6 +4,7 @@ from tinygrad.uop.ops import UOp, Ops
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.engine.realize import get_program
from tinygrad.helpers import AMX
from test.helpers import replace_opts
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
class TestFloat4(unittest.TestCase):
@@ -24,7 +25,7 @@ class TestFloat4(unittest.TestCase):
s = c.schedule()[0]
realized_ast = s.ast
opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
assert TestFloat4.count_float4(program.uops) == (2, 1)
@@ -35,8 +36,8 @@ class TestFloat4(unittest.TestCase):
c = a + b
s = c.schedule()[0]
uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]).uops
uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=2)]),
renderer=Device[Device.DEFAULT].renderer).uops
assert TestFloat4.count_float4(uops) == (4, 2)
@unittest.skipUnless(Device.DEFAULT in {"CPU"} and AMX, "Only CPU with AMX upcasts float up to size 16")
@@ -47,8 +48,8 @@ class TestFloat4(unittest.TestCase):
c = a + b
s = c.schedule()[0]
return get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]).uops
return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=shift)]),
renderer=Device[Device.DEFAULT].renderer).uops
sizes = [12, 8, 16]
shifts = [3, 2, 4]
@@ -66,7 +67,7 @@ class TestFloat4(unittest.TestCase):
s = c.schedule()[0]
realized_ast = s.ast
opts_to_apply = [Opt(op=OptOps.UPCAST, axis=0, arg=4)]
program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
program = get_program(replace_opts(realized_ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
assert TestFloat4.count_float4(program.uops) == (0, 1)
@@ -77,8 +78,8 @@ class TestFloat4(unittest.TestCase):
c = a + b
s = c.schedule()[0]
uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]).uops
uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=2)]),
renderer=Device[Device.DEFAULT].renderer).uops
assert TestFloat4.count_float4(uops) == (0, 2)
@@ -90,8 +91,8 @@ class TestFloat4(unittest.TestCase):
c = a + b
s = c.schedule()[0]
return get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
opts=[Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]).uops
return get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=shift)]),
renderer=Device[Device.DEFAULT].renderer).uops
sizes = [13, 9, 17]
shifts = [3, 2, 4]
@@ -109,7 +110,7 @@ class TestFloat4(unittest.TestCase):
# float4 should be emitted (the reduce axis of size 4 is the float4 axis here)
s = c.schedule()[0]
uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UNROLL, axis=0, arg=4)]).uops
uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UNROLL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
assert TestFloat4.count_float4(uops) == (0, 0)
@@ -123,8 +124,8 @@ class TestFloat4(unittest.TestCase):
# UPDATE: now we do this fusion
s = c.schedule()[0]
uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer,
opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
renderer=Device[Device.DEFAULT].renderer).uops
assert TestFloat4.count_float4(uops) in {(0,1), (1,1)}
@@ -137,7 +138,7 @@ class TestFloat4(unittest.TestCase):
# since the top axis is not contiguous.
s = c.schedule()[0]
uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
assert TestFloat4.count_float4(uops) == (0, 1)
@@ -149,7 +150,7 @@ class TestFloat4(unittest.TestCase):
# should float4 b but not a
s = c.schedule()[0]
uops = get_program(s.ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=4)]).uops
uops = get_program(replace_opts(s.ast, [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
assert TestFloat4.count_float4(uops) == (1, 1)
test/opt/test_tensor_cores.py
+9 -9
View File
@@ -8,7 +8,7 @@ from tinygrad.uop.ops import Ops
from tinygrad.dtype import DType
from tinygrad.device import is_dtype_supported
from tinygrad.helpers import AMX, DEV, Context
from test.helpers import slow
from test.helpers import slow, replace_opts
from tinygrad.engine.realize import CompiledRunner, get_program
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
from tinygrad.codegen.opt.tc import amd_cdna_1616128
@@ -27,14 +27,14 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
opts_to_apply = [Opt(OptOps.TC, axis, (tc_select, tc_opt, 1))]
if ensure_triggered:
program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
wmmas = len([uop for uop in program.uops if uop.op is Ops.WMMA])
tcs = len([x for x in program.applied_opts if x.op is OptOps.TC])
assert wmmas > 0, "tensor core not triggered"
assert tcs == 1, "tensor core opt not included"
else:
try:
program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts_to_apply)
program = get_program(replace_opts(realized_ast, opts_to_apply), Device[Device.DEFAULT].renderer)
assert False, "OptOps.TC triggered, expected KernelOptError"
except KernelOptError: pass
@@ -45,7 +45,7 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
if dtype_in == dtypes.bfloat16: r = r.float()
realized_ast, bufs = helper_realized_ast(r)
opts = [Opt(op=OptOps.TC, axis=axis, arg=(tc_select, tc_opt, use_tensor_cores))]
prg = CompiledRunner(replace(get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=opts), device=Device.DEFAULT))
prg = CompiledRunner(replace(get_program(replace_opts(realized_ast, opts), Device[Device.DEFAULT].renderer), device=Device.DEFAULT))
if use_tensor_cores == 1: assert len([uop for uop in prg.p.uops if uop.op is Ops.WMMA]) > 0, "wmma not triggered"
assert len([x for x in prg.p.uops[-1].arg.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
prg.exec(bufs)
@@ -74,7 +74,7 @@ class TestTensorCores(unittest.TestCase):
n, m, k = tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2]
a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
r = a.matmul(b, dtype=tc.dtype_out)
prg = get_program(r.schedule()[-1].ast, Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))])
prg = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer)
if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM":
assert "0x201000" in prg.src
elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM":
@@ -141,7 +141,7 @@ class TestTensorCores(unittest.TestCase):
c = a.conv2d(b, padding=1, dtype=tc.dtype_out)
realized_ast, real_bufs = helper_realized_ast(c)
program = get_program(realized_ast, Device[Device.DEFAULT].renderer, opts=[Opt(OptOps.TC, axis, (-1, 2, 1))])
program = get_program(replace_opts(realized_ast, [Opt(OptOps.TC, axis, (-1, 2, 1))]), Device[Device.DEFAULT].renderer)
assert len([uop for uop in program.uops if uop.op is Ops.WMMA]) > 0, "tensor core not triggered"
assert len([x for x in program.applied_opts if x.op is OptOps.TC]) == 1, "tensor core opt not included"
@@ -165,7 +165,7 @@ class TestTensorCores(unittest.TestCase):
r = x.matmul(y, dtype=tc.dtype_out)
opts = [Opt(OptOps.UNROLL, 0, 4)]
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
if u.op is Ops.WMMA:
assert u.src[-1].src[0].op != Ops.STORE
@@ -179,7 +179,7 @@ class TestTensorCores(unittest.TestCase):
r = x.matmul(y, dtype=tc.dtype_out)
opts = [Opt(OptOps.UNROLL, 0, 4)]
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
if u.op is Ops.WMMA:
#assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
assert u.src[-1].src[0].op != Ops.STORE
@@ -195,7 +195,7 @@ class TestTensorCores(unittest.TestCase):
r = x.matmul(y, dtype=tc.dtype_out).relu()
opts = [Opt(OptOps.UNROLL, 0, 4)]
ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
for u in get_program(ast, Device[Device.DEFAULT].renderer, opts=opts).uops:
for u in get_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).uops:
if u.op is Ops.WMMA:
#assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
assert u.src[-1].src[0].op != Ops.STORE
tinygrad/codegen/__init__.py
+1 -6
View File
@@ -6,7 +6,6 @@ from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_
from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
from tinygrad.renderer import Renderer, ProgramSpec, Estimates
from tinygrad.dtype import dtypes
from tinygrad.codegen.opt import Opt
# import all pattern matchers here
from tinygrad.codegen.gpudims import pm_add_gpudims
@@ -151,7 +150,7 @@ pm_to_program = PatternMatcher([
@Context(ALLOW_DEVICE_USAGE=0)
@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ast), ret=renderer), replay=True)
def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> ProgramSpec:
def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
"""
Transform an AST into a ProgramSpec. May trigger BEAM search.
@@ -168,10 +167,6 @@ def get_program(ast:UOp, renderer:Renderer, opts:list[Opt]|None=None) -> Program
beam, ast = (ast.arg, ast.src[0]) if ast.op is Ops.BEAM else (0, ast)
# rewrite to prg
assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to get_program"
if opts is not None:
# TODO: should this be here?
assert ast.arg.opts_to_apply is None, "can't apply opts if there's already opts to apply"
ast = ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts)))
full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None, beam=beam)
prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)))
else: