don't test search

This commit is contained in:
2025-09-05 15:15:05 -07:00
parent 9a6d7f9a1e
commit 5bac311df4
2 changed files with 0 additions and 85 deletions
-6
View File
@@ -20,12 +20,6 @@ repos:
language: system
always_run: true
pass_filenames: false
- id: devicetests
name: select GPU tests
entry: env GPU=1 PYTHONPATH="." python3 -m pytest test/test_uops.py test/test_search.py
language: system
always_run: true
pass_filenames: false
- id: tests
name: subset of tests
entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
-79
View File
@@ -1,79 +0,0 @@
import unittest
from tinygrad import Tensor, Device
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.codegen.opt.postrange import Scheduler
from tinygrad.codegen.opt.search import bufs_from_lin, actions, beam_search
from tinygrad.helpers import Context, GlobalCounters
from tinygrad.engine.realize import capturing
class TestBEAM(unittest.TestCase):
def test_dynamic_beam(self):
# TODO: make this infra globally usable
class Capture:
def __init__(self): self.captured = []
def add(self, x): self.captured.append(x)
capturing.append(Capture())
kernel_count = GlobalCounters.kernel_count
with Context(BEAM=1): Tensor.zeros(16).contiguous().realize()
assert GlobalCounters.kernel_count == kernel_count + 1
k_beam_1 = capturing[0].captured
capturing.clear()
capturing.append(Capture())
kernel_count = GlobalCounters.kernel_count
with Context(BEAM=0): Tensor.zeros(16).contiguous().realize()
assert GlobalCounters.kernel_count == kernel_count + 1
k_beam_0 = capturing[0].captured
capturing.clear()
self.assertNotEqual(k_beam_0[-1].prg.p.src, k_beam_1[-1].prg.p.src)
def test_get_kernel_actions_dedup(self):
from test.test_linearizer import helper_realized_ast
from tinygrad.codegen.opt.search import get_kernel_actions
a = Tensor.empty(4, 3)
b = Tensor.empty(3)
realized_ast, _ = helper_realized_ast(a @ b)
candidates = [
Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=4),
Opt(op=OptOps.LOCAL, axis=0, arg=0), Opt(op=OptOps.LOCAL, axis=0, arg=4),
Opt(op=OptOps.UNROLL, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=3),
Opt(op=OptOps.GROUP, axis=0, arg=0), Opt(op=OptOps.GROUP, axis=0, arg=3),
Opt(op=OptOps.GROUPTOP, axis=0, arg=0), Opt(op=OptOps.GROUPTOP, axis=0, arg=3),
]
lins = get_kernel_actions(Scheduler(realized_ast, Device.default.renderer), include_0=False, candidates=candidates).values()
# ensure amt=0 are not duplicated
assert all(len(x.applied_opts) == 1 for x in lins)
kernel_actions = [x.applied_opts[0] for x in lins]
assert Opt(OptOps.UPCAST, axis=0, arg=4) not in kernel_actions, "did not de-dup UPCAST"
assert Opt(OptOps.LOCAL, axis=0, arg=4) not in kernel_actions, "did not de-dup LOCAL"
assert Opt(OptOps.UNROLL, axis=0, arg=3) not in kernel_actions, "did not de-dup UNROLL"
assert Opt(OptOps.GROUP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUP"
assert Opt(OptOps.GROUPTOP, axis=0, arg=3) not in kernel_actions, "did not de-dup GROUPTOP"
def test_get_kernel_actions_preserves_actions_state(self):
from test.test_linearizer import helper_realized_ast
from tinygrad.codegen.opt.search import get_kernel_actions
a = Tensor.rand(16, 16)
b = Tensor.rand(16, 16)
realized_ast, _ = helper_realized_ast(a @ b)
actions_before = actions.copy()
get_kernel_actions(Scheduler(realized_ast))
actions_after = actions.copy()
assert actions_after == actions_before, "actions state was not preserved"
def test_beam_unnamed_kernels(self):
from test.test_linearizer import push_views
a = Tensor.rand(100)
b = Tensor.rand(100)
si = (a+b).schedule()[-1]
lin = Scheduler(push_views(si.ast), Device.default.renderer)
bufs = bufs_from_lin(lin)
# TODO: beam should have better instrumentation so we don't have to check this indirect thing
kcount = len(Scheduler.kernel_cnt)
beam_search(lin, bufs, 3, disable_cache=True)
self.assertEqual(kcount, len(Scheduler.kernel_cnt))
if __name__ == '__main__':
unittest.main()