From 80d9cced075e22defe2d948b5ecc1c05ca3d1938 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 6 Aug 2025 15:05:21 -0700 Subject: [PATCH] more test cleanups (#11544) * more test cleanups * revert that --- test/external/external_benchmark_schedule.py | 27 ++-------- test/test_arange.py | 5 +- test/unit/test_conv_shapetracker.py | 55 -------------------- 3 files changed, 6 insertions(+), 81 deletions(-) delete mode 100644 test/unit/test_conv_shapetracker.py diff --git a/test/external/external_benchmark_schedule.py b/test/external/external_benchmark_schedule.py index 2a47061e69..3dce947828 100644 --- a/test/external/external_benchmark_schedule.py +++ b/test/external/external_benchmark_schedule.py @@ -1,12 +1,8 @@ -from typing import List from extra.models.resnet import ResNet50 -from tinygrad import Tensor, nn -from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen +from tinygrad import Tensor, nn, Device +from tinygrad.helpers import Profiling, Timing, getenv from tinygrad.uop.ops import Ops -from tinygrad.opt.kernel import Kernel -from tinygrad.opt.heuristic import hand_coded_optimizations from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer -from tinygrad.opt.search import beam_search, bufs_from_lin from tinygrad.uop.spec import type_verify if __name__ == "__main__": @@ -31,26 +27,13 @@ if __name__ == "__main__": if not SCHEDULE_ONLY: asts = list({x.ast.key:x.ast for x in sched if x.ast.op is Ops.SINK}.values()) if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1] - kernels: List[Kernel] = [] - with Timing(f"***** model opts({len(asts):2d}) in "): - with Profiling(PROFILE >= 3): - for ast in asts: - k = Kernel(ast) - if BEAM: - with Context(DEBUG=max(2, DEBUG.value)): k = beam_search(k, bufs_from_lin(k), BEAM.value) - elif NOOPT: pass - else: k.apply_opts(hand_coded_optimizations(k)) - kernels.append(k) - - with Timing("***** model prep in "): - kernels = [(k, k.get_optimized_ast(), get_rewrites_for_renderer(k.opts, linearizer=False)) for k in kernels] + rewrites = get_rewrites_for_renderer(Device.default.renderer, linearizer=False) with Profiling(PROFILE, fn="/tmp/rewrite.prof"): with Timing("***** model rewrite in "): rewritten_uops = [] - for i,(k,u,rewrites) in enumerate(kernels): - with Timing(f"rewrite {i:2d} {k.name}{' '*(50-ansilen(k.name))}", enabled=getenv("VERBOSE", 0)): - rewritten_uops.append(apply_rewrites(u, rewrites)) + for u in asts: + rewritten_uops.append(apply_rewrites(u, rewrites)) if LINEARIZE: with Timing("***** model linearize in "): diff --git a/test/test_arange.py b/test/test_arange.py index bb645c77c3..6faaa41b4c 100644 --- a/test/test_arange.py +++ b/test/test_arange.py @@ -14,10 +14,7 @@ class TestArange(unittest.TestCase): tt = Tensor.arange(N) sched = tt.schedule() self.assertEqual(len(sched), 1) - k = Kernel(sched[-1].ast) - if opts is not None: - for o in opts: k.apply_opt(o) - p = get_program(k.get_optimized_ast(), k.opts) + p = get_program(sched[-1].ast, opts=opts) print(p.name) #print(p.src) ExecItem(CompiledRunner(p), [tt.uop.buffer]).run() diff --git a/test/unit/test_conv_shapetracker.py b/test/unit/test_conv_shapetracker.py deleted file mode 100644 index e812827ac1..0000000000 --- a/test/unit/test_conv_shapetracker.py +++ /dev/null @@ -1,55 +0,0 @@ -#!/usr/bin/env python -import unittest -from tinygrad.uop.ops import Ops -from tinygrad.tensor import Tensor -from tinygrad.nn import Conv2d -from tinygrad.shape.shapetracker import ShapeTracker, View -from tinygrad.helpers import prod -from test.unit.test_shapetracker import shapetracker_getitem - -class TestConvShapetracker(unittest.TestCase): - def test_conv_3x3_one_view(self): - conv = Conv2d(16, 32, (3, 3)) - # first run to init the weights, they are scheduled. - conv(Tensor.empty(1, 16, 10, 10)).schedule() - # run it again to get the kernels - sched = [si for si in conv(Tensor.empty(1, 16, 10, 10)).schedule() if si.ast.op is Ops.SINK] - assert len(sched) == 1, f"conv should only have one kernel, getting {len(sched)}" - for st in [x.st_arg for x in sched[0].ast.toposort() if x.op is Ops.LOAD]: - assert len(st.views) == 1 - - def test_conv_2x2_backward_one_view(self): - X = Tensor.rand(1, 1, 3, 3, requires_grad=True) - conv = Conv2d(1, 1, (2, 2), bias=False) - conv(X).mean().backward() - si = X.grad.schedule()[-1] - print(si) - ldb = [x for x in si.ast.toposort() if x.op is Ops.LOAD][0] - st: ShapeTracker = ldb.st_arg.simplify() - print(si.bufs[1].size) - self.assertEqual(si.bufs[1].size, st.real_size()) - for v in st.views: print(v) - - # same st - test_st = ShapeTracker(( - View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False), - View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0, - mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False))) - #test_st = ShapeTracker(( - # View(shape=(2,4), strides=(1,4), offset=0, mask=None, contiguous=False), - #)).simplify() - #View(shape=(1, 1, 2, 4, 2, 4), strides=(0, 0, 2, 8, 1, 4), offset=0, mask=((0, 1), (0, 1), (0, 2), (0, 2), (0, 2), (0, 2)), contiguous=False), - #View(shape=(1, 1, 1, 1, 3, 3, 3, 3), strides=(0, 0, 0, 0, 24, 8, 3, 1), offset=0, - # mask=((0, 1), (0, 1), (0, 1), (0, 1), (0, 2), (0, 3), (0, 2), (0, 3)), contiguous=False))).simplify() - print("*** new ***") - for v in test_st.views: print(v) - for i in range(prod(st.shape)): - i1, i2 = shapetracker_getitem(st, i), shapetracker_getitem(test_st, i) - print(i, i1, i2, si.bufs[1].size, i1==i2) - #self.assertEqual(i1, i2) - - with self.assertRaises(AssertionError): - assert len(st.views) <= 2 - -if __name__ == '__main__': - unittest.main()