more KernelCountException

* add new schedule tests + format better

* assert_kernel_count

.github/workflows/benchmark.yml

@@ -104,6 +104,9 @@ jobs:
         ./extra/amdpci/setup_python_cap.sh
         ./extra/hcq/hcq_smi.py amd rmmod
         ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -155,6 +158,9 @@ jobs:
         ./extra/amdpci/setup_python_cap.sh
         ./extra/hcq/hcq_smi.py amd rmmod
         ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -204,6 +210,9 @@ jobs:
         ./extra/amdpci/setup_python_cap.sh
         ./extra/hcq/hcq_smi.py amd rmmod
         ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: Symlink models and datasets
       run: |
         mkdir -p extra/datasets
@@ -250,6 +259,9 @@ jobs:
         ./extra/amdpci/setup_python_cap.sh
         ./extra/hcq/hcq_smi.py amd rmmod
         ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -292,6 +304,9 @@ jobs:
         ./extra/amdpci/setup_python_cap.sh
         ./extra/hcq/hcq_smi.py amd rmmod
         ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |

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
+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
@@ -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())
@@ -157,7 +153,7 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       z = emb(x).realize()
       self.assertLessEqual(GlobalCounters.global_ops, op_limit)
-      self.assertEqual(GlobalCounters.kernel_count, 2)
+      assert_kernel_count(2)
     if getenv("CHECK", 1):
       import torch
       with torch.no_grad():
@@ -257,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")

test/backend/test_custom_kernel.py

@@ -4,6 +4,7 @@ import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import assert_kernel_count
 
 # **** kernels ****
 
@@ -276,7 +277,7 @@ class TestCustomKernel(unittest.TestCase):
 
     GlobalCounters.reset()
     out.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 5)
+    assert_kernel_count(5)
 
   def test_simple_reshape(self):
     a = Tensor.ones(2,3,4).realize()
@@ -286,7 +287,7 @@ class TestCustomKernel(unittest.TestCase):
     GlobalCounters.reset()
     c.realize()
     assert all(i == 3. for i in c.flatten().tolist()), f"all 3 {c.tolist()}"
-    self.assertEqual(GlobalCounters.kernel_count, 3)
+    assert_kernel_count(3)
 
   def test_multi_after_schedule_order(self):
     """Test correct scheduling order when custom_kernel has multiple outputs.
@@ -336,7 +337,7 @@ class TestCustomKernel(unittest.TestCase):
     c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
     GlobalCounters.reset()
     c.realize()
-    self.assertEqual(GlobalCounters.kernel_count, len(devs))
+    assert_kernel_count(len(devs))
     self.assertTrue((c == 2).all().item())
 
   def test_partial_invalid_store_keeps_uncovered_reads(self):
@@ -401,7 +402,7 @@ class TestCustomKernel(unittest.TestCase):
     else: z = y.T.T+1
     GlobalCounters.reset()
     z.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    assert_kernel_count(2)
     self.assertEqual(z.tolist(), x.add(2).tolist())
 
   @unittest.expectedFailure
@@ -418,7 +419,7 @@ class TestCustomKernel(unittest.TestCase):
     GlobalCounters.reset()
     y = run(x[0]).realize()
     # it's copying the input and the output
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(y.tolist(), [1, 2, 3, 4])
 
   @Context(DEV="CPU")

test/backend/test_jit.py

@@ -2,7 +2,7 @@
 import unittest
 import numpy as np
 
-from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
+from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, Variable, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
@@ -97,7 +97,7 @@ class TestJit(unittest.TestCase):
       prev = o
 
     # Checking that 2 graphs are inited.
-    assert len(jf.captured.linear.src) == 2
+    if len(jf.captured.linear.src) != 2: raise KernelCountException(2, len(jf.captured.linear.src))
     for si in jf.captured.linear.src:
       assert call_is_graph(si)
 

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)

test/backend/test_llama_kernels.py

@@ -7,7 +7,7 @@ from extra.llama_kernels import local_abs_max
 from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
-from test.helpers import needs_second_gpu
+from test.helpers import needs_second_gpu, assert_kernel_count
 from test.backend.test_asm_gemm import has_hipcc
 
 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
@@ -95,7 +95,7 @@ class TestLocalAmax(unittest.TestCase):
     x = Tensor.arange(16).reshape(4, 4).cast(dtypes.float).clone(devices[0]).realize().shard(devices, axis=0).realize()
     GlobalCounters.reset()
     out = (x * local_abs_max(x)).clone().realize()
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    assert_kernel_count(2)
     self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
 
 @unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")

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, assert_kernel_count
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -62,7 +62,7 @@ class TestMultiTensor(unittest.TestCase):
   def test_shard_empty(self):
     GlobalCounters.reset()
     X = Tensor.empty(256).shard(devices_2, 0).realize()
-    assert GlobalCounters.kernel_count == 0
+    assert_kernel_count(0)
     (X + X).realize()
 
   # TODO: fix this to not copy on the src device
@@ -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)
 

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)

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

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)
 

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, assert_kernel_count
 
 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)
 
@@ -126,9 +103,9 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.arange(16).clone().realize()
     GlobalCounters.reset()
     a[4] = 3
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     a.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertListEqual(a.tolist(), [0, 1, 2, 3, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
 
   def test_no_extra_contiguous_on_setitem_assign_back(self):

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):

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

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, GlobalCounters
 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,36 @@ def derandomize_model(model):
     p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
     p.realize()
 
+class KernelCountException(Exception):
+  def __init__(self, expected:int, got:int):
+    self.expected, self.got = expected, got
+    super().__init__(f"expected {expected}, got {got}")
+
+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(allowed, kernel_cnt)
+  # test compiling the linear
+  compile_linear(linear)
+  return linear, var_vals
+
+def assert_kernel_count(expected:int):
+  got = GlobalCounters.kernel_count
+  if got != expected: raise KernelCountException(expected, got)
+
 def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]
   return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"
@@ -53,15 +84,15 @@ def jit_cache_count(linear:UOp) -> int:
 def assert_jit_cache_len(fxn, expected_len):
   linear = fxn.captured.linear if fxn.captured is not None else None
   if linear is None or not linear.src:
-    assert expected_len == 0, expected_len
+    if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
   if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
   if call_is_graph(linear.src[0]):
-    assert len(linear.src) == 1, len(linear.src)
+    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION
-    assert len(inner.src) == expected_len, f"expected {expected_len}, got {len(inner.src)}"
+    if len(inner.src) != expected_len: raise KernelCountException(expected_len, len(inner.src))
   else:
-    assert len(linear.src) == expected_len, f"expected {expected_len}, got {len(linear.src)}"
+    if len(linear.src) != expected_len: raise KernelCountException(expected_len, len(linear.src))
 
 def min_normal(dt:DType) -> float: return 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))
 

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)

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(max_kernels_allowed, kernels_used)
+      if (max_kernels_allowed - kernels_used) / max_kernels_allowed >= 0.2:
+        raise KernelCountException(max_kernels_allowed, kernels_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
 

test/null/test_schedule.py

@@ -2,32 +2,11 @@
 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.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 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, KernelCountException
 
 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)
@@ -224,7 +202,7 @@ class TestSchedule(unittest.TestCase):
     GlobalCounters.reset()
     expr = (a/b)/c
     expr.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertLessEqual(GlobalCounters.global_ops, 4*3)
 
   # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
@@ -357,6 +335,11 @@ class TestSchedule(unittest.TestCase):
     out1 = a.sum() + b
     check_schedule([out0, out1], 2)
 
+  def test_reduce_broadcast_not_recomputed(self):
+    a = Tensor.empty(32, 16).realize()
+    out = a-a.mean(axis=0, keepdim=True)
+    check_schedule(out, 2)
+
   def test_scaled_dot_product_attention_multireduce_fusion(self):
     q = Tensor.empty(32,8,16,8).realize()
     k = Tensor.empty(32,8,16,8).realize()
@@ -609,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)
 
@@ -632,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):
@@ -647,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):
@@ -714,6 +697,19 @@ class TestSchedule(unittest.TestCase):
     xt = X[[Tensor([2]), Tensor([1])]]
     check_schedule(xt, 1)
 
+  def test_split_advanced_indexing_not_recomputed(self):
+    with Context(SPLIT_REDUCEOP=1):
+      X = Tensor.empty(32768, 4).realize()
+      idx = Tensor.randint(4, high=X.shape[0])
+      linear, _ = check_schedule(X[idx], 3, [Tensor._device_rng_counters[idx.device]])
+      # The split's final reduction remains, but the one-hot gather should collapse into a direct indexed load.
+      reduce_kernels = 0
+      for call in linear.src:
+        if call.src[0].op is not Ops.SINK: continue
+        sink = full_rewrite_to_sink(call.src[0], Device[call.device].renderer)
+        reduce_kernels += any(u.op is Ops.RANGE and u.arg[-1] is AxisType.REDUCE for u in sink.toposort())
+      self.assertEqual(reduce_kernels, 1)
+
   def test_push_through_reshape(self):
     x = Tensor.empty(10, 20).realize()
     out = x.argmax(1)
@@ -874,8 +870,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 +1452,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 +1867,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 +1876,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

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)

test/null/test_uops.py

@@ -55,11 +55,11 @@ class TestDTypeFromUOp(unittest.TestCase):
     invalid = UOp.invalid()
     self.assertIs(invalid.dtype, dtypes.bool)
     self.assertIs(UOp.const(Invalid, dtypes.float32), invalid)
-    self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved)
     scratch = Tensor.invalids(4, dtype=dtypes.float32)
     self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort()
       if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
     invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
+    for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
     for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
               UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
               UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)

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

test/unit/test_assign.py

@@ -4,6 +4,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
+from test.helpers import assert_kernel_count
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -42,7 +43,7 @@ class TestAssign(unittest.TestCase):
     # it should copy into the empty buffer
     GlobalCounters.reset()
     c.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -50,7 +51,7 @@ class TestAssign(unittest.TestCase):
     xs.assign(xs+1)
     GlobalCounters.reset()
     self.assertListEqual(X.tolist(), [1,2,4,5])
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
 
   def test_assign_slice_alt(self):
     X = Tensor([1,2,3,4]).realize()
@@ -58,7 +59,7 @@ class TestAssign(unittest.TestCase):
     xs1.assign(xs2+1)
     GlobalCounters.reset()
     self.assertListEqual(X.tolist(), [1,4,5,4])
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    assert_kernel_count(2)
 
   def test_assign_flip(self):
     ref = np.arange(16, dtype=np.float32)
@@ -68,7 +69,7 @@ class TestAssign(unittest.TestCase):
     xs.assign(xs + X)
     ref = ref + ref[::-1]
     np.testing.assert_allclose(X.numpy(), ref)
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    assert_kernel_count(2)
 
   def test_assign_add(self):
     for T in (1, 2, 10):#, 100): # this crashes in CI, not sure why
@@ -331,14 +332,14 @@ class TestAssign(unittest.TestCase):
     a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1)
     GlobalCounters.reset()
     b.assign(a.contiguous()).realize()
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    assert_kernel_count(2)
 
   def test_assign_contiguous_permute(self):
     b = Tensor.arange(16).reshape(4,4).clone().realize()
     a = (Tensor.arange(16).reshape(4,4).clone().realize() + 1).permute((1,0))
     GlobalCounters.reset()
     b.assign(a.contiguous()).realize()
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    assert_kernel_count(2)
 
   def test_permuted_assignment(self):
     a = Tensor(np.arange(N*N, dtype=np.float32)).reshape(N,N)
@@ -413,7 +414,7 @@ class TestAssign(unittest.TestCase):
 
     GlobalCounters.reset()
     Tensor.realize(b, c, d)
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     np.testing.assert_allclose(b.numpy(), a.sum(1).numpy()+1)
     np.testing.assert_allclose(c.numpy(), a.sum(1).numpy()+2)
     np.testing.assert_allclose(d.numpy(), a.sum(1).numpy()+3)
@@ -461,7 +462,7 @@ class TestAssign(unittest.TestCase):
     b.assign(r + b)
     c.assign(r + b_perm.contiguous())
     Tensor.realize(b, c)
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    assert_kernel_count(2)
     np.testing.assert_equal(b.numpy(), a.numpy().sum(1) + np.arange(32 * 32).reshape(32, 32))
     np.testing.assert_equal(c.numpy(), a.numpy().sum(1) + np.arange(32 * 32).reshape(32, 32).transpose(1, 0))
 
@@ -471,7 +472,7 @@ class TestAssign(unittest.TestCase):
     a.assign(a + b)
     GlobalCounters.reset()
     a.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     np.testing.assert_equal(a.numpy(), np.ones((4, 4))+np.pad(np.ones((4, 4))[:, 0:2], ((0, 0), (0, 2)), constant_values=2))
 
   def test_permuted_assignment_masked_view_not_contiguous(self):
@@ -510,7 +511,7 @@ class TestAssign(unittest.TestCase):
     expected[0:10] = expected[50:60].copy()
     GlobalCounters.reset()
     a[0:10].assign(a[50:60]).realize()
-    self.assertEqual(GlobalCounters.kernel_count, 2)  # currently conservative, forces contiguous
+    assert_kernel_count(2)  # currently conservative, forces contiguous
     np.testing.assert_allclose(a.numpy(), expected)
 
   def test_setitem_half(self):
@@ -630,7 +631,7 @@ class TestAssign(unittest.TestCase):
     GlobalCounters.reset()
     x.realize()
     # N assigns (1 kernel each) producing N kernels total
-    self.assertEqual(GlobalCounters.kernel_count, N)
+    assert_kernel_count(N)
 
   def test_shared_computation_assign_kernel_count(self):
     """When a .contiguous() is shared between an assign value and the next layer's input (like QKV projection in LLM),
@@ -648,7 +649,7 @@ class TestAssign(unittest.TestCase):
     GlobalCounters.reset()
     caches[-1][:1].contiguous().realize()
     # N matmuls + N assigns + 1 final read = 2*N+1 (AFTER embedding allows full graph scheduling with shared contiguous reuse)
-    self.assertEqual(GlobalCounters.kernel_count, 2*N+1)
+    assert_kernel_count(2*N+1)
 
   def test_double_assign_from_const(self):
     a = Tensor.empty(2)
@@ -656,7 +657,7 @@ class TestAssign(unittest.TestCase):
     a.assign(Tensor.ones(2, buffer=False))
     GlobalCounters.reset()
     a.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(a.tolist(), [1.,1.])
 
   def test_assign_deviceless_const(self):
@@ -672,7 +673,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    self.assertEqual(GlobalCounters.kernel_count, 2)  # TODO: first copy is dead, could be 1
+    assert_kernel_count(2)  # TODO: first copy is dead, could be 1
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -682,7 +683,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(base.tolist(), [1,4,3])
 
 class TestAssignOrdering(unittest.TestCase):

test/unit/test_dtype_weak.py

@@ -44,9 +44,13 @@ class TestWeakPromotion(unittest.TestCase):
     r = Tensor([2], dtype=dtypes.uint8, device="CPU").copysign(Tensor([1], dtype=dtypes.uint32, device="CPU"))
     self.assertEqual((r.dtype, r.tolist()), (dtypes.uint32, [2]))
 
-  def test_minimum_commits_both_operands(self):
+  def test_minimum_reflects_weak_operand(self):
     r = Tensor(1).minimum(Tensor([2], dtype=dtypes.uint8, device="CPU"))
     self.assertEqual((r.dtype, r.tolist()), (dtypes.uint8, [1]))
+    for dt in dtypes.uints:
+      r = Tensor([dt.max], dtype=dt, device="CPU").minimum(1)
+      self.assertEqual((r.dtype, r.tolist()), (dt, [1]))
+      self.assertNotIn(Ops.CAST, [u.op for u in r._uop.toposort()])
 
   def test_broadcasted_keeps_const_weak(self):
     # a python scalar stays a bare weak CONST through _broadcasted, lifted only to the KIND of the lub

test/unit/test_function.py

@@ -4,6 +4,7 @@ from tinygrad.function import function
 from tinygrad import Tensor, GlobalCounters, Device
 from tinygrad.dtype import Invalid
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
+from test.helpers import assert_kernel_count
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -618,7 +619,7 @@ class TestFunctionTuple(unittest.TestCase):
     out = f(a)
     GlobalCounters.reset()
     out.realize()
-    self.assertEqual(GlobalCounters.kernel_count, kernel_count)
+    assert_kernel_count(kernel_count)
     np.testing.assert_allclose(out.numpy(), [3., 5., 7., 9.])
 
   def test_custom_kernel_precompile_further_compute_multi(self): self.test_custom_kernel_precompile_further_compute(multi=True, kernel_count=4)

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
+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))
@@ -143,9 +143,8 @@ class TestMultiTensor(unittest.TestCase):
     GlobalCounters.reset()
     with Context(ALLREDUCE_CAST=1, RING=0, ALL2ALL=0):
       tst.realize()
-    kernel_count = GlobalCounters.kernel_count
+    assert_kernel_count(kernel_count)
     np.testing.assert_allclose(tst.numpy(), (a_src.numpy()+b_src.numpy()).sum(0))
-    self.assertEqual(kernel_count, kernel_count)
 
   def test_allreduce_cast_half_assign(self): self.test_allreduce_cast_half(assign=True, kernel_count=10)
 
@@ -584,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")

test/unit/test_setitem_schedule.py

@@ -1,15 +1,16 @@
 import unittest
 from tinygrad import Tensor, dtypes, GlobalCounters
+from test.helpers import assert_kernel_count
 
 class TestSetitemInto(unittest.TestCase):
   def test_setitem_into_unrealized(self):
     GlobalCounters.reset()
     t = Tensor.arange(4, dtype=dtypes.int32).reshape(2, 2)
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     self.assertEqual(GlobalCounters.global_mem, 0)
     self.assertListEqual(t.tolist(), [[0, 1], [5, 5]])
 
@@ -18,11 +19,11 @@ class TestSetitemInto(unittest.TestCase):
     GlobalCounters.reset()
     a = Tensor.arange(8, dtype=dtypes.int32).reshape(2, 4)
     w = a[0] + a[1]  # unrealized ADD with SHRINK in graph: [4, 6, 8, 10]
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     w[1] = 99
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     w.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     self.assertEqual(GlobalCounters.global_mem, 0)
     self.assertListEqual(w.tolist(), [4, 99, 8, 10])
 
@@ -30,61 +31,61 @@ class TestSetitemInto(unittest.TestCase):
     GlobalCounters.reset()
     t = Tensor.empty(4, dtype=dtypes.int32)
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(GlobalCounters.global_mem, 4)
     t[1].realize()
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(t[1].item(), 5)
 
   def test_setitem_into_empty_alu(self):
     GlobalCounters.reset()
     t = Tensor.empty(4, dtype=dtypes.int32) + 1
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertLessEqual(GlobalCounters.global_mem, 32)
     t[1].realize()
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(t[1].item(), 5)
 
   def test_setitem_into_tensor(self):
     t = Tensor([1, 2, 3, 4], dtype=dtypes.int32).realize()
     GlobalCounters.reset()
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t[1].realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(GlobalCounters.global_mem, 4)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertListEqual(t.tolist(), [1, 5, 3, 4])
 
   def test_setitem_into_tensor_alu(self):
     t = Tensor([1, 2, 3, 4], dtype=dtypes.int32).realize() + 1
     GlobalCounters.reset()
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t[1].realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertLessEqual(GlobalCounters.global_mem, 32)
     t[1].realize()
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertListEqual(t.tolist(), [2, 5, 4, 5])
 
   def test_setitem_into_const(self):
     GlobalCounters.reset()
     t = Tensor.ones(4, dtype=dtypes.int32, buffer=False)
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     self.assertEqual(GlobalCounters.global_mem, 0)
     self.assertListEqual(t.tolist(), [1, 5, 1, 1])
 
@@ -92,9 +93,9 @@ class TestSetitemInto(unittest.TestCase):
     GlobalCounters.reset()
     t = Tensor.ones(4, dtype=dtypes.int32, buffer=False) + 1
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     self.assertEqual(GlobalCounters.global_mem, 0)
     self.assertListEqual(t.tolist(), [2, 5, 2, 2])
 
@@ -105,18 +106,18 @@ class TestSetitemInto(unittest.TestCase):
     t = Tensor.arange(4, dtype=dtypes.int32)
     self.assertIs(other.uop, t.uop)
     t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     self.assertListEqual(t.tolist(), [0, 5, 2, 3])
 
   def test_setitem_slice_const(self):
     t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
     GlobalCounters.reset()
     t[20:50] = 3
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(GlobalCounters.global_mem, 30*4)  # 30 elements written
 
   def test_setitem_slice_tensor(self):
@@ -124,18 +125,18 @@ class TestSetitemInto(unittest.TestCase):
     v = Tensor.zeros(30, dtype=dtypes.int32).contiguous().realize()
     GlobalCounters.reset()
     t[20:50] = v
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(GlobalCounters.global_mem, 30*4*2)  # 30 read + 30 written
 
   def test_setitem_full(self):
     t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
     GlobalCounters.reset()
     t[:] = 3
-    self.assertEqual(GlobalCounters.kernel_count, 0)
+    assert_kernel_count(0)
     t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    assert_kernel_count(1)
     self.assertEqual(GlobalCounters.global_mem, 100*4)  # full buffer written
 
 if __name__ == '__main__':

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:

tinygrad/mixin/dtype.py

@@ -30,7 +30,7 @@ class DTypeMixin:
     print(t.dtype, t.numpy())
     ```
     """
-    return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.is_invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
+    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
 
   def bitcast(self, dtype:DTypeLike) -> Self:
     """

tinygrad/mixin/elementwise.py

@@ -24,6 +24,7 @@ class ElementwiseMixin(CreationMixin):
     out_dtype = least_upper_dtype(x.dtype, y.dtype)
     # keep weak CONST weak, might lift weakint -> weakfloat
     def promote(t):
+      if t._uop.base.is_invalid: return t  # invalid bool is weak const
       if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype)))
       return t.cast(out_dtype)
     return promote(x), promote(y)
@@ -395,9 +396,10 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     t, x = self._broadcasted(x)
-    # ~ is width-dependent: min(a,b) == ~max(~a,~b) only holds at a common width, so a weak operand commits at its sibling's
-    t, x = t.cast(dt:=least_upper_dtype(t.dtype, x.dtype)), x.cast(dt)
-    return t._inverse().maximum(x._inverse())._inverse()
+    # NOTE: the int inverse is done in python, since const has weak dtype without width
+    # TODO: clean this up once _broadcasted does not promote dtype
+    if dtypes.is_float(dt:=least_upper_dtype(t.dtype, x.dtype)): return -(-t).alu(Ops.MAX, -x)
+    return (t ^ (k:=dt.const(dt.min+dt.max))).alu(Ops.MAX, x ^ k) ^ k
 
   def copysign(self, other: Self | ConstType) -> Self:
     """

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<y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)))),
   (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),

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

tinygrad/uop/ops.py

@@ -801,7 +801,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
         # arg is the other srcs; all are cast to the promoted dtype, spec requires STACK srcs to match its dtype
         srcs = (self,)+tuple(arg)
         dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None))
-        return UOp(Ops.STACK, dtype, tuple(u.cast(dtype) for u in srcs))
+        # TODO: why cast here?
+        return UOp(Ops.STACK, dtype, tuple(u if u.base.is_invalid else u.cast(dtype) for u in srcs))
       case _: raise RuntimeError(f"{op} is not a MovementOp")
     usrcs = [shape_to_shape_arg(arg) for arg in src_args]
     if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg)
@@ -1096,11 +1097,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
@@ -1135,6 +1139,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:
@@ -1758,7 +1763,7 @@ def lower_weak_node(u:UOp) -> UOp|None:
   if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
   dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
                     else unwrap(dtype_from_uop(u.op, src, u.arg)))
-  return u.replace(dtype=None, src=src[:start]+tuple(s.cast(dt) for s in src[start:])).cast(u.dtype)
+  return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else s.cast(dt) for s in src[start:])).cast(u.dtype)
 pm_lower_weak = PatternMatcher([
   (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
   # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.

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 []))+")"),