From 46f0003776edf427f2984923e737a8312f6ba70b Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:55:38 -0700 Subject: [PATCH] more KernelCountException (#17407) --- test/backend/test_arange.py | 6 +++--- test/null/test_schedule.py | 12 ++++++------ test/unit/test_multitensor.py | 4 ++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/test/backend/test_arange.py b/test/backend/test_arange.py index 16e7d22e29..d58e86e5f2 100644 --- a/test/backend/test_arange.py +++ b/test/backend/test_arange.py @@ -4,13 +4,13 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable from tinygrad.helpers import Context, getenv, DEV from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear from tinygrad.renderer.ptx import PTXRenderer -from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count +from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException class TestArange(unittest.TestCase): def _get_flops(self, tensor, desired): GlobalCounters.reset() linear = compile_linear(tensor.schedule_linear()) - self.assertEqual(len(linear.src), 1) + if len(linear.src) != 1: raise KernelCountException(1, len(linear.src)) run_linear(linear) np.testing.assert_equal(tensor.numpy(), desired) return estimate_uop(linear.src[-1]).ops @@ -253,7 +253,7 @@ class TestIndexing(unittest.TestCase): xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis) xq_rope.sum().backward() linear = compile_linear(wq.grad.schedule_linear()) - assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}" + if len(linear.src) != 1: raise KernelCountException(1, len(linear.src)) bwd_ops = estimate_uop(linear.src[0]).ops expected_ops = bs*seqlen*dim*dim*ops_scale print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops") diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index f0a5b5c4d3..bcd9487c4f 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -6,7 +6,7 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo, AxisType from tinygrad.helpers import GlobalCounters, Context from tinygrad.engine.realize import run_linear, compile_linear from tinygrad.codegen import to_program, full_rewrite_to_sink -from test.helpers import check_schedule, assert_kernel_count +from test.helpers import check_schedule, assert_kernel_count, KernelCountException def _realize_weights(m): for p in nn.state.get_parameters(m): p.realize() @@ -592,9 +592,7 @@ class TestSchedule(unittest.TestCase): img = Tensor.randn(BS, CIN, 64, 64).realize() w = Tensor.uniform(16, CIN, 3, 3).realize() ret = Tensor.conv2d(img, w).relu().mean().backward() - linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad) - cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK]) - assert cnt == allowed, f"expected {allowed} kernels, got {cnt}" + check_schedule([ret, img.grad, w.grad], allowed) def test_conv2d_half(self): self.test_conv2d(4, dtype=dtypes.half) @@ -615,7 +613,8 @@ class TestSchedule(unittest.TestCase): return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM]) with Context(IMAGE=1): - self.assertEqual(cnt(), 5) + got = cnt() + if got != 5: raise KernelCountException(5, got) def test_image_f16_residual_fusion(self): with Context(FLOAT16=1, OPENPILOT_HACKS=1): @@ -630,7 +629,8 @@ class TestSchedule(unittest.TestCase): return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM]) with Context(IMAGE=1): - self.assertEqual(cnt(), 9) + got = cnt() + if got != 9: raise KernelCountException(9, got) def _test_fusion(self, shapes, f, cnt): with Context(DEBUG=0, TRACK_MATCH_STATS=0): diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index 96b90ff345..2eb12e8db4 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -2,7 +2,7 @@ import unittest, numpy as np from tinygrad import Tensor, Variable, Context, Device, TinyJit, GlobalCounters, dtypes, UOp, nn, getenv from tinygrad.nn.state import get_parameters, get_state_dict from tinygrad.uop.ops import Ops -from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count +from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count, KernelCountException from hypothesis import given, strategies as strat, settings settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) @@ -583,7 +583,7 @@ class TestMultiTensor(unittest.TestCase): zeros = Tensor.zeros(3).realize() b = a.to(devices_2)*zeros.to(devices_2) sched = b.schedule_linear().src - self.assertEqual(len(sched), 0) + if len(sched) != 0: raise KernelCountException(0, len(sched)) self.assertListEqual(b.tolist(), [0, 0, 0]) @unittest.skipIf(not_support_multi_device(), "no multi")