From 7b6d2ddf23b02a3e4b625432e1900525016da8cc Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 4 Aug 2026 15:10:43 -0400 Subject: [PATCH 1/5] more weak const without cast in const_like [PR] (#17395) * more weak const without cast in const_like [PR] * that? --- tinygrad/codegen/__init__.py | 2 +- tinygrad/uop/symbolic.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 3f37a4c1dd..58fceb0882 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -213,7 +213,7 @@ def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp): topo = r.src[0].toposort() ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END]) input_ranges = tuple(x for x in topo if x.op is Ops.RANGE and x not in r.src[1:] and x not in ended_ranges) - acc_init = acc.after(*input_ranges).store(identity_element(r.arg[0], r.dtype)) + acc_init = acc.after(*input_ranges).store(UOp.const(identity_element(r.arg[0], r.dtype))) acc_initted = acc.after(acc_init, *r.src[1:]) inp = r.src[0].reduce(arg=r.arg) if r.arg[1] else r.src[0] acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable") diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 21b1b622e3..65bc354c26 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -13,7 +13,7 @@ from tinygrad.codegen.decomp.transcendental import xpow # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** def simplify_pow(x:UOp, c:UOp) -> UOp|None: - if c.val < 0: return x.reciprocal().pow(-c) + if c.val < 0: return x.reciprocal().pow(-c.val) if c.val == 0: return x.const_like(1) if int(c.val-0.5)+0.5 == c.val: return x.pow(c.val-0.5) * x.sqrt() if int(c.val) == c.val: return (y := x.pow(c.val//2)) * y * (x if c.val%2 == 1 else 1) From c1a10e07264ef168856bb3368681ccdc512fa871 Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 4 Aug 2026 18:06:50 -0400 Subject: [PATCH 2/5] fix _min_max for CAST from float to int [pr] (#17396) * fix _min_max for CAST from float to int [pr] * fix --- test/null/test_uop_vmin_vmax.py | 6 ++++++ tinygrad/uop/ops.py | 9 ++++++--- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index fb7a6bee4d..fd5c3a7ab2 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -162,6 +162,12 @@ class TestVminVmaxProperties(unittest.TestCase): self.assertEqual(x_uint.vmin, dtypes.uint.min) self.assertEqual(x_uint.vmax, dtypes.uint.max) + def test_vmin_vmax_cast_float_to_int(self): + self.assertEqual(UOp.variable('x', -4.5, 4.5, dtypes.float).cast(dtypes.int)._min_max, (-4, 4)) + self.assertEqual(UOp.const(4.5).cast(dtypes.float).cast(dtypes.int)._min_max, (4, 4)) + x = UOp.const(4.5).cast(dtypes.float) + self.assertIs(x.ne(x.cast(dtypes.int).cast(dtypes.float)).simplify().arg, True) + def test_vmin_vmax_invalid(self): i = UOp.invalid() self.assertNotEqual(i.vmin, i.vmax) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index dc19b168ea..fac3b5d2ad 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1096,11 +1096,14 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.CONST and self.val is not Invalid: return self.val, self.val if self.op is Ops.INDEX: return self.src[0]._min_max if self.op is Ops.CAST: + # an int destination truncates a float source toward zero. trunc is monotone + smin, smax = self.src[0]._min_max + if dtypes.is_int(self.dtype) and dtypes.is_float(self.src[0].dtype) and all(math.isfinite(v) for v in (smin, smax)): + smin, smax = math.trunc(smin), math.trunc(smax) # a cast to unsigned keeps exact bounds when the source fits # TODO: can do more based on new dtype window - if dtypes.is_unsigned(self.dtype) and 0 <= self.src[0].vmin and self.src[0].vmax <= self.dtype.max: return self.src[0]._min_max - if self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): - return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max) + if dtypes.is_unsigned(self.dtype) and 0 <= smin and smax <= self.dtype.max: return smin, smax + if self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): return max(self.dtype.min, smin), min(smax, self.dtype.max) return self.dtype.min, self.dtype.max @functools.cached_property From d79772f0575c9432aca88e73094fefcc1a44e451 Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 4 Aug 2026 19:30:43 -0400 Subject: [PATCH 3/5] fix pow on extreme inputs (#17397) * fix pow on extreme inputs * WEBGPU --- test/backend/test_ops.py | 11 +++++++++++ tinygrad/codegen/decomp/transcendental.py | 8 ++++---- tinygrad/mixin/gradient.py | 2 +- 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/test/backend/test_ops.py b/test/backend/test_ops.py index 520ed9d793..ae2c81ee15 100644 --- a/test/backend/test_ops.py +++ b/test/backend/test_ops.py @@ -728,6 +728,17 @@ class TestOps(unittest.TestCase): else: self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}") + def test_pow_neg_inf_frac_exponent(self): + # pow(-inf, 0.3) is +inf, so the gradient 0.3*pow(-inf, -0.7) is 0, never nan + helper_test_op(None, lambda x: x**0.3, vals=[[-math.inf]]) + # is_odd truncates, so it calls 3.3 odd: only the non_int guard keeps pow(-inf, 3.3) from negating to -inf + helper_test_op(None, lambda x: x**3.3, vals=[[-math.inf]]) + + def test_pow_zero_exponent(self): + # x ** 0 is the constant 1 for every x, so the gradient with respect to the base is 0, never nan + # TODO: nan ** 0, failed on WEBGPU + helper_test_op(None, lambda x,y: x**y, vals=[[-math.inf, math.inf, 0.0], [0.0, 0.0, 0.0]]) + def test_pow_zero_tensor(self): helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.0]]) # TODO: fix WEBGPU diff --git a/tinygrad/codegen/decomp/transcendental.py b/tinygrad/codegen/decomp/transcendental.py index b2771acb37..e4e66fbd5a 100644 --- a/tinygrad/codegen/decomp/transcendental.py +++ b/tinygrad/codegen/decomp/transcendental.py @@ -257,12 +257,12 @@ def xlog2(d:UOp) -> UOp: def xpow(base:UOp, exponent:UOp) -> UOp: # start with b ** e = exp2(e * log2(b)) ret = (base < 0).where(-base, base).log2().mul(exponent).exp2() - # negative base: nan for non-integer exponent, negate for odd integer exponent + # negative base: nan for non-integer exponent, negate for odd integer exponent. -inf is never nan, it stays |base| ** exponent non_int = exponent != exponent.cast(dtypes.int32).cast(exponent.dtype) is_odd = (exponent < 0).where(-exponent, exponent).cast(dtypes.int32).mod(2).cast(dtypes.bool) - neg_base = non_int.where(ret.const_like(math.nan), is_odd.where(-ret, ret)) - # fix 0 ** 0 = 1 - return (base.eq(0) & exponent.eq(0)).where(ret.const_like(1), (base < 0).where(neg_base, ret)) + neg_base = non_int.where(base.ne(-math.inf).where(ret.const_like(math.nan), ret), is_odd.where(-ret, ret)) + # x ** 0 = 1, including 0 ** 0 and inf ** 0 + return exponent.eq(0).where(ret.const_like(1), (base < 0).where(neg_base, ret)) @functools.cache def get_transcendental_patterns(ops:tuple[Ops, ...], force_transcendental:bool) -> PatternMatcher: diff --git a/tinygrad/mixin/gradient.py b/tinygrad/mixin/gradient.py index 57d5e0dc41..93cc63843b 100644 --- a/tinygrad/mixin/gradient.py +++ b/tinygrad/mixin/gradient.py @@ -53,7 +53,7 @@ pm_gradient = PatternMatcher([ (UPat((Ops.CMPLT, Ops.CMPNE)), lambda: (None, None)), (UPat(Ops.ADD), lambda ctx: (ctx, ctx)), (UPat(Ops.POW, name="ret", src=(UPat.var("b"), UPat.var("e"))), lambda ctx, ret, b, e: - (ctx * (b.eq(0)&e.eq(0)).where(e, e*b.pow(e-1)), ctx * b.eq(0).where((e<0).where(ret.const_like(-math.inf), 0), ret*b.log2()*math.log(2.0)))), + (ctx * e.eq(0).where(e, e*b.pow(e-1)), ctx * b.eq(0).where((e<0).where(ret.const_like(-math.inf), 0), ret*b.log2()*math.log(2.0)))), (UPat(Ops.MAX, src=(UPat.var("x"), UPat.var("y"))), lambda ctx, x, y: ((x>y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)), (x Date: Tue, 4 Aug 2026 19:44:34 -0400 Subject: [PATCH 4/5] use weak 0 in convert_pad_to_where_to_keep_behavior_local [pr] (#17398) --- tinygrad/schedule/indexing.py | 2 +- tinygrad/uop/ops.py | 1 + tinygrad/uop/render.py | 4 +++- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 04a5e67c23..b3b6d6da68 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -101,7 +101,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp): if x not in ctx.range_map: return None bx = create_bufferize_and_index_based_on_ranges(ctx, x) valid: UOp = UOp.const(True).uprod([r.get_valid() for r in ctx.range_map[x][0]]) - return valid.where(bx.src[0], UOp.const(0, x.dtype)) + return valid.where(bx.src[0], UOp.const(x.dtype.const(0))) def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp): if x.arg[1] == 0: return None diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index fac3b5d2ad..4f8eabd954 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1138,6 +1138,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @staticmethod def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, device=None, volatile=False): + dtype = strong_dtype(dtype) # storage is never weak: a placeholder commits the width of what's put in it if addrspace is AddrSpace.GLOBAL: ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace, device=device,volatile=volatile)) else: diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index c8467ae654..03133e1cea 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -104,8 +104,10 @@ pm_pyrender_extra = PatternMatcher([ (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"), # `.where` re-promotes its operands, so render WHERE via .alu() too (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.WHERE, {ctx[x.src[1]]}, {ctx[x.src[2]]})"), + # the binary operators re-promote their operands (a weak src meeting a strong one gets a cast), render those via .alu() too (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x: - strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")), + strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})") + if x.src[0]._broadcasted(x.src[1]) == x.src else f"{ctx[x.src[0]]}.alu({x.op}, {ctx[x.src[1]]})"), (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"), (UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \ ([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"), From 6122b3c98fbc1bbc840c95e095bbab9e9150d8eb Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:17:29 -0700 Subject: [PATCH 5/5] use check_schedule in tests where possible (#17400) --- test/backend/test_arange.py | 14 +++----- test/backend/test_linearizer.py | 5 ++- test/backend/test_multitensor.py | 5 ++- test/backend/test_nn.py | 10 ++---- test/backend/test_renderer_failures.py | 4 +-- test/backend/test_schedule.py | 33 +++---------------- test/backend/test_softmax_fusion.py | 4 +-- test/backend/test_transcendental.py | 2 +- test/helpers.py | 25 ++++++++++++++- test/null/test_attention.py | 5 ++- test/null/test_real_world.py | 7 ++-- test/null/test_schedule.py | 44 ++++++-------------------- test/null/test_winograd.py | 3 +- 13 files changed, 63 insertions(+), 98 deletions(-) diff --git a/test/backend/test_arange.py b/test/backend/test_arange.py index b2faa47ef9..676116f431 100644 --- a/test/backend/test_arange.py +++ b/test/backend/test_arange.py @@ -4,7 +4,7 @@ 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 +from test.helpers import needs_second_gpu, check_schedule class TestArange(unittest.TestCase): def _get_flops(self, tensor, desired): @@ -55,8 +55,7 @@ class TestIndexing(unittest.TestCase): with Context(NOOPT=1): GlobalCounters.reset() out = ((Tensor.arange(1,16385)-1)*needle).sum() - linear, var_vals = out.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(out, 1) run_linear(linear, var_vals) self.assertEqual(out.item(), 1337) @@ -72,8 +71,7 @@ class TestIndexing(unittest.TestCase): reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1) full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False)) X = full.sum(axis=(2,3)) - linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(X, 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}" np.testing.assert_allclose(real_index, X.numpy()) @@ -98,8 +96,7 @@ class TestIndexing(unittest.TestCase): GlobalCounters.reset() X = dataset[idxs] assert X.shape == (4,DDIM) - linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(X, 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}" np.testing.assert_allclose(real_index, X.numpy()) @@ -113,8 +110,7 @@ class TestIndexing(unittest.TestCase): GlobalCounters.reset() X = dataset[idxs] assert X.shape == (4,DDIM) - linear, var_vals = X.linear_with_vars() - self.assertEqual(len(linear.src), 1) + linear, var_vals = check_schedule(X, 1) run_linear(linear, var_vals) assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}" np.testing.assert_allclose(real_index, X.numpy()) diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 6d833cbeba..958a9be840 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -12,7 +12,7 @@ from tinygrad.dtype import DType, dtypes, AddrSpace from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.cstyle import CUDARenderer from tinygrad.renderer.isa import ISARenderer -from test.helpers import replace_opts +from test.helpers import replace_opts, check_schedule from test.backend.test_softmax_fusion import single_kernel_softmax MOCKGPU = DEV.interface.startswith("MOCK") @@ -293,8 +293,7 @@ class TestLinearizer(unittest.TestCase): a = Tensor.ones(4, 4).contiguous().realize() b = a.shrink(((1, 2), None)).pad(((1, 2), None)).bool() a.assign(b.where(2, a)) - linear, var_vals = a.linear_with_vars() - assert len(linear.src) == 1 + linear, var_vals = check_schedule(a, 1) run_linear(linear, var_vals) np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.]) program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer) diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index 5e461236c1..7ef0a6888b 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -6,7 +6,7 @@ from tinygrad.nn.state import get_parameters from tinygrad.engine.realize import run_linear, compile_linear import numpy as np from hypothesis import given, strategies as strat, settings -from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph +from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False)) settings.load_profile("my_profile") @@ -355,8 +355,7 @@ class TestMultiTensor(unittest.TestCase): def test_const_like_shrink_on_shard_axis(self): t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0) out = t.const_like(2)[:, :8] - linear, var_vals = out.linear_with_vars() - self.assertEqual(len(linear.src), 0) + linear, var_vals = check_schedule(out, 0) run_linear(linear, var_vals) self.assertEqual(out.tolist(), [[2]*8]*16) diff --git a/test/backend/test_nn.py b/test/backend/test_nn.py index 88c4174c4f..75195d88e0 100644 --- a/test/backend/test_nn.py +++ b/test/backend/test_nn.py @@ -3,11 +3,11 @@ import unittest import numpy as np import torch from tinygrad import Tensor, Device, TinyJit, dtypes -from tinygrad.uop.ops import Ops from tinygrad.helpers import GlobalCounters, Context from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell from tinygrad.nn.state import load_state_dict +from test.helpers import check_schedule from tinygrad.engine.realize import run_linear from test.helpers import not_support_multi_device, needs_second_gpu, slow @@ -428,18 +428,14 @@ class TestNN(unittest.TestCase): a = Tensor([[1, 5, 9, 11], [12, 19, 8, 1]]) result = layer(a) - linear, var_vals = result.linear_with_vars() - self.assertEqual(len([call for call in linear.src if call.src[0].op is Ops.SINK]), kcount, - "first run realizes weight and embedding") + linear, var_vals = check_schedule(result, kcount) run_linear(linear, var_vals) b = Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) result = layer(b) - linear, var_vals = result.linear_with_vars() - self.assertEqual(1, len([call for call in linear.src if call.src[0].op is Ops.SINK]), - "second run realizes embedding only") + linear, var_vals = check_schedule(result, 1) run_linear(linear, var_vals) print(f"Embedding used {GlobalCounters.global_ops} ops") self.assertLessEqual(GlobalCounters.global_ops, ops) diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index 826928bece..d1f7a145fa 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -8,6 +8,7 @@ from tinygrad.helpers import prod from tinygrad.renderer.cstyle import CStyleLanguage from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.wgsl import WGSLRenderer +from test.helpers import check_schedule from tinygrad.runtime.ops_python import PythonRenderer from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu from tinygrad.tensor import Tensor @@ -61,8 +62,7 @@ class TestCStyleFailures(unittest.TestCase): dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None ret = Tensor.empty(1, dtype=dtype) for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype)) - linear = ret.schedule_linear() - assert len(linear.src) == 1 + linear, _ = check_schedule(ret, 1) src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[2].arg self.assertEqual("("*5 not in src, should_strip_paren) diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index 838b8cbcdd..2d86674633 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -6,34 +6,13 @@ import unittest, time import numpy as np from tinygrad import nn, dtypes, Device, Tensor, Variable -from tinygrad.uop.ops import UOp, Ops, UPat -from tinygrad.helpers import DEBUG, DEV, GlobalCounters, Context, all_same, temp -from tinygrad.engine.realize import compile_linear, run_linear +from tinygrad.uop.ops import Ops, UPat +from tinygrad.helpers import DEV, GlobalCounters, Context, all_same, temp +from tinygrad.engine.realize import run_linear +from test.helpers import check_schedule supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes() -class KernelCountException(Exception): pass -def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): - if to_prerealize: - with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize) - if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars() - elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t) - else: - assert isinstance(t, UOp), f"can't schedule {t}" - linear, var_vals = Tensor(t).linear_with_vars() - kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) - for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) - if kernel_cnt != allowed: - print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") - if DEBUG >= 3: - for i,call in enumerate(linear.src): - print("kernel", i+1) - print(call.src[0]) - raise KernelCountException(f"{kernel_cnt} != {allowed}") - # test compiling the linear - compile_linear(linear) - return linear, var_vals - def _realize_weights(m): for p in nn.state.get_parameters(m): p.realize() @@ -113,11 +92,9 @@ class TestSchedule(unittest.TestCase): a2 = mop(a) expected = (a+a2).tolist() a.assign(a+a2) - linear, var_vals = a.linear_with_vars() - kcount = len(linear.src) + linear, var_vals = check_schedule(a, expected_kcount) run_linear(linear, var_vals) self.assertListEqual(a.tolist(), expected) - self.assertEqual(kcount, expected_kcount) def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2) def test_setitem_paddded_sched(self): self.test_setitem_sched(lambda x: x.shrink_to(4, 1).pad_to(4, 4), 1) diff --git a/test/backend/test_softmax_fusion.py b/test/backend/test_softmax_fusion.py index eb0dda3ade..97f76139ee 100644 --- a/test/backend/test_softmax_fusion.py +++ b/test/backend/test_softmax_fusion.py @@ -4,6 +4,7 @@ from tinygrad import Tensor, GlobalCounters, Context, Device from tinygrad.dtype import DTypeLike, dtypes from tinygrad.engine.realize import run_linear from tinygrad.helpers import DEBUG, get_single_element +from test.helpers import check_schedule def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor: # only support axis =-1 @@ -103,8 +104,7 @@ class TestFuse(unittest.TestCase): k = (x @ wk).contiguous() v = (x @ wv).contiguous() attn = q.scaled_dot_product_attention(k, v) - s = attn.schedule_linear() - self.assertEqual(len(s.src), 4) # 3 matmul and 1 attention + check_schedule(attn, 4) # 3 matmul and 1 attention @unittest.skip("needs RANGEIFY>1") def test_flash_attention(self): diff --git a/test/backend/test_transcendental.py b/test/backend/test_transcendental.py index fb96ad5147..20f008154d 100644 --- a/test/backend/test_transcendental.py +++ b/test/backend/test_transcendental.py @@ -2,7 +2,7 @@ import unittest from tinygrad import Tensor, Device, dtypes from tinygrad.tensor import _to_np_dtype from tinygrad.helpers import Context, getenv, DEV, OSX -from test.backend.test_schedule import check_schedule +from test.helpers import check_schedule from test.backend.test_dtype_alu import ht, dtypes_float import numpy as np import math diff --git a/test/helpers.py b/test/helpers.py index 1e2b893d16..c613d78f1f 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -8,10 +8,11 @@ from tinygrad.tensor import _to_np_dtype from tinygrad.codegen import to_program from tinygrad.dtype import DType, truncate from tinygrad.nn.state import get_parameters -from tinygrad.helpers import T, Target, DEV +from tinygrad.helpers import T, Target, DEV, DEBUG, Context from tinygrad.renderer import Renderer from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups from tinygrad.codegen.late.linearizer import linearize +from tinygrad.engine.realize import compile_linear # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run") @@ -34,6 +35,28 @@ def derandomize_model(model): p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype)) p.realize() +class KernelCountException(Exception): pass +def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): + if to_prerealize: + with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize) + if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars() + elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t) + else: + assert isinstance(t, UOp), f"can't schedule {t}" + linear, var_vals = Tensor(t).linear_with_vars() + kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) + for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) + if kernel_cnt != allowed: + print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") + if DEBUG >= 3: + for i,call in enumerate(linear.src): + print("kernel", i+1) + print(call.src[0]) + raise KernelCountException(f"{kernel_cnt} != {allowed}") + # test compiling the linear + compile_linear(linear) + return linear, var_vals + def call_is_graph(call:UOp) -> bool: ast = call.src[0] return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph" diff --git a/test/null/test_attention.py b/test/null/test_attention.py index aa23608de9..2c03e58450 100644 --- a/test/null/test_attention.py +++ b/test/null/test_attention.py @@ -1,7 +1,7 @@ import unittest from tinygrad import Tensor, dtypes, TinyJit, UOp from tinygrad.llm.model import apply_rope as apply_rope_new, precompute_freqs_cis -from test.helpers import assert_jit_cache_len +from test.helpers import assert_jit_cache_len, check_schedule def apply_rope(x:Tensor, start_pos:int): B, H, T, Hd = x.shape @@ -16,9 +16,8 @@ class TestAttention(unittest.TestCase): k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() attn = q.scaled_dot_product_attention(k, v) - sched = attn.schedule_linear() # attention has 4 kernels now - self.assertEqual(len(sched.src), 4) + check_schedule(attn, 4) def test_apply_rope_jit_prune(self): def rope_fn(x_in, pos): return apply_rope(x_in, pos) diff --git a/test/null/test_real_world.py b/test/null/test_real_world.py index 378e8b924b..23dcab609f 100644 --- a/test/null/test_real_world.py +++ b/test/null/test_real_world.py @@ -5,7 +5,7 @@ from tinygrad.nn.state import get_parameters from tinygrad.engine.jit import TinyJit from tinygrad import Tensor, Device, GlobalCounters, dtypes, Variable from tinygrad.helpers import Context -from test.helpers import slow, jit_cache_count +from test.helpers import slow, jit_cache_count, KernelCountException from extra.lr_scheduler import OneCycleLR from test.helpers import derandomize_model @@ -35,8 +35,9 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used" assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used" if kernels_used: - assert kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}" - assert (max_kernels_allowed - kernels_used) / max_kernels_allowed < 0.2, f"{max_kernels_allowed=} is too far from {kernels_used=} used" + if kernels_used > max_kernels_allowed: raise KernelCountException(f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}") + if (max_kernels_allowed - kernels_used) / max_kernels_allowed >= 0.2: + raise KernelCountException(f"{max_kernels_allowed=} is too far from {kernels_used=} used") if all_jitted: assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted" # noqa: E501 diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index 2f4cbae50a..29dd939281 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -3,31 +3,10 @@ import gc, unittest, time from typing import cast from tinygrad import nn, dtypes, Device, Tensor, getenv from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo -from tinygrad.helpers import DEBUG, GlobalCounters, Context -from tinygrad.engine.realize import compile_linear, run_linear +from tinygrad.helpers import GlobalCounters, Context +from tinygrad.engine.realize import run_linear, compile_linear from tinygrad.codegen import to_program - -class KernelCountException(Exception): pass -def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): - if to_prerealize: - with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize) - if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars() - elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t) - else: - assert isinstance(t, UOp), f"can't schedule {t}" - linear, var_vals = Tensor(t).linear_with_vars() - kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1) - for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink) - if kernel_cnt != allowed: - print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}") - if DEBUG >= 3: - for i,call in enumerate(linear.src): - print("kernel", i+1) - print(call.src[0]) - raise KernelCountException(f"{kernel_cnt} != {allowed}") - # test compiling the linear - compile_linear(linear) - return linear, var_vals +from test.helpers import check_schedule def _realize_weights(m): for p in nn.state.get_parameters(m): p.realize() @@ -143,7 +122,7 @@ class TestSimpleSchedule(unittest.TestCase): a = Tensor.empty(16,16).sum(axis=1) a1 = a.reshape(4,4) a2 = a.reshape(16,1,1) - self.assertEqual(len(Tensor.schedule_linear(a1, a2).src), 1) + check_schedule([a1, a2], 1) class TestSchedule(unittest.TestCase): def setUp(self): @@ -155,8 +134,7 @@ class TestSchedule(unittest.TestCase): def test_arange_avgpool2d(self, kcount=1): x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32) t = x.avg_pool2d(padding=1).clone() - linear, var_vals = t.linear_with_vars() - self.assertEqual(len(linear.src), kcount) + check_schedule(t, kcount) def test_arange_avgpool2d_fused_noopt(self): with Context(NOOPT=1): self.test_arange_avgpool2d(kcount=1) @@ -874,8 +852,7 @@ class TestSchedule(unittest.TestCase): t = Tensor.zeros((3, 3)).contiguous().realize() v = t[1] # view - is_realized but not has_buffer_identity assert v.uop.is_realized - linear, _ = Tensor.linear_with_vars(v) - self.assertEqual(len(linear.src), 0) + check_schedule(v, 0) # NOTE: because empty does not have a lowered kernel if realize is called on a childless empty, it never gets allocated. def test_childless_empty_never_allocates(self): @@ -1457,8 +1434,7 @@ class TestSchedule(unittest.TestCase): Tensor.manual_seed(0) x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize() out = x.softmax(dtype=dtypes.float) - linear = out.schedule_linear() - self.assertEqual(len(linear.src), 3) + linear, _ = check_schedule(out, 3) # max reduction stays in input dtype (no numerical loss), upcast happens after subtracting max self.assertEqual(linear.src[0].src[1].dtype, dtypes.half) self.assertEqual(linear.src[1].src[1].dtype, dtypes.float) @@ -1873,8 +1849,7 @@ class TestFusionOp(unittest.TestCase): val = 1.0 a = Tensor(val) for _ in range(24): a = Tensor.stack(a, a)[0] - linear = a.schedule_linear() - self.assertLessEqual(len(linear.src), 1) + check_schedule(a, 0) self.assertLess(time.perf_counter()-st, 2.0) def test_recursive_reshape(self): @@ -1883,8 +1858,7 @@ class TestFusionOp(unittest.TestCase): b = Tensor.empty(16, 2).realize() r = a.sum(1) for _ in range(24): r = r.reshape(16, 2) + b - linear = r.schedule_linear() - self.assertEqual(len(linear.src), 1) + check_schedule(r, 1) self.assertLess(time.perf_counter()-st, 2.0) # NOTE: the NULL backend supports SLICE diff --git a/test/null/test_winograd.py b/test/null/test_winograd.py index 1530675828..0fa0e5039e 100644 --- a/test/null/test_winograd.py +++ b/test/null/test_winograd.py @@ -1,6 +1,7 @@ import unittest, sys from tinygrad import Tensor, GlobalCounters, dtypes, Context from tinygrad.helpers import WINO +from test.helpers import check_schedule @unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows") class TestWinograd(unittest.TestCase): @@ -13,7 +14,7 @@ class TestWinograd(unittest.TestCase): def test_forward_kernels(self): x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize() out = Tensor.conv2d(x,w) - self.assertEqual(len(out.schedule_linear().src), 4) + check_schedule(out, 4) def test_backward_counters(self): # contiguous_backward on the pooled input keeps the input-transform adjoint out of the overlap accumulation, so