From 9ef319f3493f661fb742fcccda90577e2037d71a Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 08:56:22 +0800 Subject: [PATCH 01/29] bad conv in rangeify (#12373) * bad conv with broken rangeify * no maxpool needed * add empty_like * typo * no self * issue remains for test --- test/test_rangeify.py | 7 +++++++ tinygrad/tensor.py | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/test/test_rangeify.py b/test/test_rangeify.py index a7cb8a4a54..d15b993965 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -38,6 +38,13 @@ class TestRangeifyOpt(unittest.TestCase): Xsel, Ysel = X[sel], Y[sel] Tensor.realize(Xsel, Ysel) + def test_resnetconv(self): + conv1 = nn.Conv2d(3, 8, kernel_size=7, stride=2, bias=False, padding=3) + conv1.weight.replace(conv1.weight.empty_like()) + x = Tensor.empty(1, 3, 56, 56) + x = conv1(x).pad([1,1,1,1])+1 + x.realize() + @unittest.skipIf(RANGEIFY<1, "tests only for RANGEIFY") class TestRangeify(unittest.TestCase): def test_groupnorm(self): diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 81ef5013e7..76969ea11d 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -456,6 +456,13 @@ class Tensor(MathTrait): device = tuple(Device.canonicalize(d) for d in device) if isinstance(device, tuple) else Device.canonicalize(device) return Tensor(UOp.new_buffer(device, size, dtype), device, dtype, **kwargs).shrink(((0,prod(shape)),)).reshape(shape) + def empty_like(self, **kwargs) -> Tensor: + """ + Creates an empty tensor with the same shape as `self`. + If `dtype` is not specified, the dtype of `self` is used. + """ + return Tensor.empty(self.shape, dtype=kwargs.pop("dtype", self.dtype), device=kwargs.pop("device", self.device), **kwargs) + @staticmethod def from_blob(ptr:int, shape:tuple[int, ...], **kwargs) -> Tensor: """ From 969a1b35ca3934810b099de31ecbc29e440cca59 Mon Sep 17 00:00:00 2001 From: hooved <172129504+hooved@users.noreply.github.com> Date: Tue, 30 Sep 2025 21:21:08 -0400 Subject: [PATCH 02/29] LR scheduler for Stable Diffusion mlperf training (#12201) * add lr scheduler for stable diffusion training * add lr scheduler test * rerun ci * rerun CI * use np for testing * move test to CI path * remove unneeded copy --- examples/mlperf/lr_schedulers.py | 25 +++++++++++++++++++++++-- test/external/external_test_optim.py | 24 +++++++++++++++++++++++- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/examples/mlperf/lr_schedulers.py b/examples/mlperf/lr_schedulers.py index e339ae2d51..b8e9ffa8e5 100644 --- a/examples/mlperf/lr_schedulers.py +++ b/examples/mlperf/lr_schedulers.py @@ -1,8 +1,9 @@ import math -from tinygrad import dtypes +from tinygrad import dtypes, Tensor from tinygrad.nn.optim import Optimizer from extra.lr_scheduler import LR_Scheduler +from typing import Callable # https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py class PolynomialDecayWithWarmup(LR_Scheduler): @@ -36,4 +37,24 @@ class CosineAnnealingLRWithWarmup(LR_Scheduler): def get_lr(self): warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos()) - return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype) \ No newline at end of file + return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype) + +# Reference: https://github.com/mlcommons/training/blob/64b14a9abc74e08779a175abca7d291f8c957632/stable_diffusion/ldm/lr_scheduler.py, Lines 36-97 +class LambdaLinearScheduler: + def __init__(self, warm_up_steps:int, f_min:float, f_max:float, f_start:float, cycle_lengths:int): + self.lr_warm_up_steps, self.f_min, self.f_max, self.f_start, self.cycle_lengths = warm_up_steps, f_min, f_max, f_start, cycle_lengths + + def schedule(self, n:Tensor) -> Tensor: + warm_up = (n < self.lr_warm_up_steps) + f_warm_up = (self.f_max - self.f_start) / self.lr_warm_up_steps * n + self.f_start + return warm_up.where(f_warm_up, self.f_min + (self.f_max - self.f_min) * (self.cycle_lengths - n) / (self.cycle_lengths)) + +# based on torch.optim.lr_scheduler.LambdaLR +class LambdaLR(LR_Scheduler): + def __init__(self, optimizer:Optimizer, base_lr:Tensor, lr_lambda:Callable): + super().__init__(optimizer) + self.base_lr, self.lr_lambda = base_lr, lr_lambda + self.step() + + def get_lr(self): + return self.base_lr * self.lr_lambda(self.epoch_counter - 1) \ No newline at end of file diff --git a/test/external/external_test_optim.py b/test/external/external_test_optim.py index 622a0645e2..014601bae7 100644 --- a/test/external/external_test_optim.py +++ b/test/external/external_test_optim.py @@ -11,7 +11,7 @@ from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, AdamW from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer -from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup, CosineAnnealingLRWithWarmup +from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup, CosineAnnealingLRWithWarmup, LambdaLR, LambdaLinearScheduler from test.external.mlperf_resnet.lars_util import PolynomialDecayWithWarmup as PolynomialDecayWithWarmup_tf np.random.seed(1337) @@ -192,5 +192,27 @@ class TestCosineAnnealingLRWithWarmup(unittest.TestCase): def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20) def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100) +class TestLambdaLRLinearWarmup(unittest.TestCase): + def test_linear_lr_warmup(self): + BS, BASE_LR = 304, 2.5e-7 + lr = BS * BASE_LR + # Use a dummy Tensor parameter for optimizer because the lr_scheduler only needs the optimizer's device and lr, the params aren't touched. + optimizer = AdamW([Tensor([1.])]) + lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule + lr_scheduler = LambdaLR(optimizer, Tensor(lr, device=optimizer.device), lambda_lr_callback) + lrs = {} + + # with above settings, optimizer.lr should warm up to lr over 1000 steps linearly + for i in range(1200): + lr_scheduler.step() + if i in {0, 499, 998, 999, 1000, 1199}: + lrs[i] = optimizer.lr.item() + + np.testing.assert_allclose(lr, lrs[999], rtol=0, atol=1e-11) + np.testing.assert_equal(lrs[999], lrs[1000]) + np.testing.assert_equal(lrs[999], lrs[1199]) + np.testing.assert_allclose(lrs[999] / lrs[0], 1000, rtol=0, atol=1) + np.testing.assert_allclose(lrs[999] / lrs[499], 2, rtol=0, atol=1e-5) + if __name__ == '__main__': unittest.main() From a128fa0f8adfbbca6fb3d023f06db09092ecd4b7 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 09:25:35 +0800 Subject: [PATCH 03/29] removing double reshapes was wrong (#12375) --- tinygrad/schedule/rangeify.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 9a5ff6de56..5ed6dc7114 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -16,18 +16,7 @@ ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL} -double_reshape = PatternMatcher([ - # RESHAPE on RESHAPE is the second reshape - (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE),), name="x"), - lambda x: x.replace(src=(x.src[0].src[0],), tag=((x.src[0].tag or ())+(x.tag or ())) or None)), -]) - -earliest_rewrites = double_reshape+PatternMatcher([ - # non shape changing RESHAPE is NOOP - #(UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0] if x.src[0].shape == x.arg else None), - # DETACH and CONTIGUOUS_BACKWARD are NOOPs here, so is FUSE - #(UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0].f(Ops.NOOP, tag=x.tag)), - +earliest_rewrites = PatternMatcher([ # just removing it works... (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]), From f2eb92948d011ab5f3f24c25d43de26485a4ace0 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 1 Oct 2025 04:37:52 +0300 Subject: [PATCH 04/29] rangeify: ban view pushing (#12371) * rangeify: ban view pushing * don't shape INDEX * fix the codegen cache * make space --- tinygrad/codegen/__init__.py | 9 +++++---- tinygrad/uop/ops.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 01b5155572..962c04f698 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -1,7 +1,7 @@ from typing import Any, Callable import functools from dataclasses import dataclass -from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL +from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL, RANGEIFY from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype from tinygrad.uop.spec import type_verify from tinygrad.renderer import Renderer @@ -46,16 +46,17 @@ rewrites_for_linearizer = [ def get_rewrites_for_renderer(opts:Renderer, optimize:bool=True, linearizer:bool=True) -> list[RewriteStep]: # cache with the values of the context vars - return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value) + return _get_rewrites_for_renderer(opts, optimize, linearizer, QUANTIZE.value, DEVECTORIZE.value, TRANSCENDENTAL.value, RANGEIFY.value) @functools.cache -def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]: +def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL, + _RANGEIFY) -> list[RewriteStep]: # ** lowerer (rewrite_shapetracker_with_index) ** ret: list[RewriteStep] = [] if optimize: # view pushing - ret.extend(rewrites_for_views) + if not _RANGEIFY: ret.extend(rewrites_for_views) # lowerer first if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize")) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index cd611ffc28..a96c83b186 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -156,7 +156,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @functools.cached_property def st(self) -> ShapeTracker|None: if self.op is Ops.INDEX and self.src[0].op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.MSTACK, - Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}: + Ops.MSELECT, Ops.BUFFER, Ops.BUFFERIZE, Ops.VECTORIZE, Ops.STORE}: return None if self.op is Ops.BARRIER: return None if self.op in GroupOp.Block: return None From 26247573e1cf096568d82042d9697167afb5d414 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 1 Oct 2025 04:53:04 +0300 Subject: [PATCH 05/29] rangeify multi tests on gpu (#12376) * rangeify multi tests on gpu * fix limit_bufs --- .github/workflows/test.yml | 2 +- tinygrad/schedule/rangeify.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9d97541df2..0ea138e2c3 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -586,7 +586,7 @@ jobs: - name: some unit tests run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/unit/test_winograd.py test/unit/test_linalg.py --durations=20 - name: Test METAL=1 RANGEIFY=1 - run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py --durations=20 + run: METAL=1 RANGEIFY=1 python -m pytest -n=auto test/test_ops.py test/test_multitensor.py --durations=20 - name: Run process replay tests uses: ./.github/actions/process-replay diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 5ed6dc7114..1e53c596e3 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -463,7 +463,7 @@ def limit_bufs(ctx:RangeifyContext, root:UOp): bufs: set[UOp] = set() def gate_input(u:UOp): # TODO: add cache to fix n^2 - if is_load:=(u.op in {Ops.BUFFERIZE, Ops.BUFFER, Ops.DEFINE_VAR}): bufs.add(u) + if is_load:=(u.op in {Ops.BUFFERIZE, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_VAR}): bufs.add(u) return not is_load root.toposort(gate=gate_input) From 4c9a930de286f256f46065be4b07eb3e4de9a2ca Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 09:59:19 +0800 Subject: [PATCH 06/29] rangeify attn tests (#12377) --- test/unit/test_attention.py | 7 +++++-- tinygrad/apps/llm.py | 3 ++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/test/unit/test_attention.py b/test/unit/test_attention.py index e6a3c487bd..5043f7335a 100644 --- a/test/unit/test_attention.py +++ b/test/unit/test_attention.py @@ -2,9 +2,11 @@ import unittest from tinygrad import Tensor, dtypes, TinyJit, UOp from tinygrad.helpers import RANGEIFY from tinygrad.apps.llm import apply_rope +#from tinygrad.engine.realize import run_schedule # TODO: test_scheduler, but just in uint class TestAttention(unittest.TestCase): + @unittest.skipIf(RANGEIFY > 0, "not half on rangeify") def test_half_qkv_buffers(self): BS, seqlen, dim = 10, 4, 100 q = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() @@ -12,11 +14,12 @@ class TestAttention(unittest.TestCase): v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize() attn = q.scaled_dot_product_attention(k, v) sched = attn.schedule() + #run_schedule(sched[:]) # attention has 5 kernels now self.assertEqual(len(sched), 4 if RANGEIFY else 5) softmax_inputs = sched[1:4] - for si in softmax_inputs: - assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=}" + for i,si in enumerate(softmax_inputs): + assert all(b.dtype == dtypes.half for b in si.bufs), f"non half {si.bufs=} in kernel {i}" def test_apply_rope(self): x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32) diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index 62801117bd..50649ffe4e 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -58,7 +58,8 @@ def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor: assert (Hd & 1) == 0, "RoPE requires an even head dimension" half = Hd // 2 angles = (Tensor.arange(T, dtype="float32") + start_pos)[:, None] * (base ** (-(Tensor.arange(half, dtype="float32") / half)))[None, :] - cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype), angles.sin().reshape(1, 1, T, half).cast(x.dtype) + # contiguous here allows RoPE to be pruned in the JIT + cos, sin = angles.cos().reshape(1, 1, T, half).cast(x.dtype).contiguous(), angles.sin().reshape(1, 1, T, half).cast(x.dtype).contiguous() x_pairs = x.reshape(B, H, T, half, 2) return Tensor.stack(x_pairs[..., 0] * cos - x_pairs[..., 1] * sin, x_pairs[..., 0] * sin + x_pairs[..., 1] * cos, dim=-1).reshape(B, H, T, Hd) From 8def8145e4fc20e0ed261263754f073cb6df9fce Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 1 Oct 2025 10:58:59 +0800 Subject: [PATCH 07/29] ALLOWED_KERNEL_COUNT openpilot 0.9.4 with RANGEIFY (#12381) --- .github/workflows/test.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 0ea138e2c3..04bcb7c9d5 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -376,9 +376,10 @@ jobs: llvm: 'true' - name: Test openpilot model kernel count and gate usage run: | - ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx + ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2160 ALLOWED_GATED_READ_IMAGE=16 RANGEIFY=0 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot model with rangeify - run: RANGEIFY=1 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx + run: | + ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=33 RANGEIFY=1 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot alt model correctness (float32) run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx - name: Test openpilot fastvits model correctness (float32) From 4204edc60b3c14df29c6a9ab79fd3f52f2a05da3 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Tue, 30 Sep 2025 20:07:39 -0700 Subject: [PATCH 08/29] feat: skip test_long (#12383) --- test/unit/test_hashing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/unit/test_hashing.py b/test/unit/test_hashing.py index 1ab969b2fd..a73e1929a6 100644 --- a/test/unit/test_hashing.py +++ b/test/unit/test_hashing.py @@ -2,6 +2,7 @@ from typing_extensions import Callable import hashlib, random, unittest from tinygrad import Tensor, Device, getenv, dtypes from tinygrad.device import is_dtype_supported +from tinygrad.helpers import CI @unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64") @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI") @@ -60,6 +61,7 @@ class TestKeccak(unittest.TestCase): # self.assertEqual(bytes(Tensor(b"a" * 1000000).keccak().tolist()), # bytearray.fromhex("5c8875ae474a3634 ba4fd55ec85bffd6 61f32aca75c6d699 d0cdcb6c115891c1")) + @unittest.skipIf(CI, "times out in ci") def test_long(self): data = b"\x00" * 4 self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16)) From 1c1b4d14e963b8d3efbf1dd33c536de05bbd9034 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 11:19:48 +0800 Subject: [PATCH 09/29] minor cleaups in rangeify (#12382) * minor cleaups in rangeify * op_in_parents * don't use toposort * Revert "don't use toposort" This reverts commit 257d8e252901998e9ec6e510c41c7c01a5938edc. --- tinygrad/schedule/rangeify.py | 13 +++++++------ tinygrad/uop/ops.py | 2 ++ 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 1e53c596e3..08ca540a0f 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -138,7 +138,7 @@ def map_reshape(idx:UOp, r:UOp): for s in r.src[0].shape[::-1]: ret.append(mish % s) # NOTE: simplify will turn this to CONST mish //= s - tret = ret[0].sink(*ret[1:]).simplify().src[::-1] if len(ret) else () + tret = UOp.sink(*ret[::-1]).simplify().src return r.src[0].index(*tret, dtype=idx.dtype, arg=idx.arg) def map_pad(idx:UOp, r:UOp): @@ -295,8 +295,8 @@ def children_gate(ctx:RangeifyContext, idx:UOp, c:UOp): def might_end_axis(idx:UOp): if idx.arg is None: return None # TODO: write a proper cost function here - if all(x.op not in {Ops.BUFFER, Ops.REALIZE, Ops.BUFFERIZE} for x in idx.toposort()): return None - if all(x.op not in {Ops.REDUCE_AXIS} for x in idx.toposort()): return None + if not idx.op_in_parents(Ops.BUFFER, Ops.REALIZE, Ops.BUFFERIZE): return None + if not idx.op_in_parents(Ops.REDUCE_AXIS): return None to_end_axis = [] for i,a in enumerate(idx.src[1:]): # in RANGEIFY=1, always realize @@ -386,7 +386,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp): # if we return None, the bufferize is kept accessed_buffers = [] - def red_gate(x): + def red_gate(x:UOp): if x.op is Ops.INDEX: accessed_buffers.append(x) return False @@ -489,7 +489,6 @@ pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary) def bufferize_to_store(x:UOp): rngs = x.src[1:] shape = tuple([int(r.vmax+1) for r in rngs]) - sym_shape = tuple([ssimplify(r.src[0]) for r in rngs]) size = prod(shape) assert size > 0, f"no zero sized buffers {shape}" @@ -514,7 +513,9 @@ def bufferize_to_store(x:UOp): ret = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=x.dtype) ret = ret.forced_reshape(shape) # TODO: is this right? what if it's offset - if shape is not sym_shape: ret = ret.shrink(tuple([(0,x) for x in sym_shape])) + if any(r.src[0].op is not Ops.RANGE for r in rngs): + sym_shape = tuple([ssimplify(r.src[0]) for r in rngs]) + ret = ret.shrink(tuple([(0,x) for x in sym_shape])) return ret.replace(tag=x.tag) # handle locals diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index a96c83b186..60372d4ea3 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -134,6 +134,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass): else: ret[node] = None # second time i'm seeing this node, add it to returned toposort return ret + def op_in_parents(self, *ops:Ops): return any(x.op in ops for x in self.toposort()) + # returns map of UOps to their children in the graph rooted by self def get_children_map(self) -> dict[UOp, dict[UOp, None]]: ret: dict[UOp, dict[UOp, None]] = {} From da52006bde6a2cef71896cce93fd6feb70768d5d Mon Sep 17 00:00:00 2001 From: b1tg <33436708+b1tg@users.noreply.github.com> Date: Wed, 1 Oct 2025 11:26:36 +0800 Subject: [PATCH 10/29] rangeify: fix test_scatter_reduce (#12380) * rangeify: fix test_scatter_reduce * ext_vector_type * set alignment=1 on boolean --- .github/workflows/test.yml | 2 ++ tinygrad/renderer/cstyle.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 04bcb7c9d5..5b7cc33fa6 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -535,6 +535,8 @@ jobs: test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_tensor_variable.py \ test/test_outerworld_range.py test/test_randomness.py test/test_nn.py test/test_arange.py test/test_tensor.py test/test_optim.py \ test/test_setitem.py test/test_assign.py test/test_multitensor.py + - name: Test CPU=1 DEVECTORIZE=0 (RANGEIFY=1) + run: CPU=1 CPU_LLVM=0 RANGEIFY=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure" - name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1 run: | CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 09d25c0828..3377e17905 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -215,8 +215,8 @@ class ClangRenderer(CStyleLanguage): kernel_typedef = "__attribute__((ms_abi)) void" def render_vector_prefix(self, dt:DType) -> str: # round (down) to power of two (this is actually the default clang behavior) - alignment = 2**int(math.log2(dt.itemsize)) if getenv("ALIGNED", 1) else 1 - return f"typedef {self.render_dtype(dt.scalar())} {self.render_dtype(dt)} __attribute__((aligned({alignment}),vector_size({dt.itemsize})));" + alignment = 2**int(math.log2(dt.itemsize)) if getenv("ALIGNED", 1) and not dtypes.is_bool(dt) else 1 + return f"typedef {self.render_dtype(dt.scalar())} {self.render_dtype(dt)} __attribute__((aligned({alignment}),ext_vector_type({dt.count})));" def _render_defines(self, uops) -> list[str]: prefix = [self.render_vector_prefix(dt) for dt in uops_to_dtypes(uops) if dt.count > 1] From 0662946fac84758bfa09434685d16f289add73e7 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 1 Oct 2025 12:05:47 +0800 Subject: [PATCH 11/29] atol in test_two_binops_no_rerun (#12387) for RANGEIFY LLVM --- test/unit/test_conv.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/unit/test_conv.py b/test/unit/test_conv.py index 6091a9c46c..47a24fa7f7 100644 --- a/test/unit/test_conv.py +++ b/test/unit/test_conv.py @@ -51,8 +51,8 @@ class TestConv(unittest.TestCase): w = Tensor.randn(32,12,3,3) out = x.conv2d(w, stride=(2,2), padding=(1,1)) r1, r2 = out.relu(), (out-1) - np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0)) - np.testing.assert_allclose(r2.numpy(), out.numpy() - 1) + np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-6) + np.testing.assert_allclose(r2.numpy(), out.numpy() - 1, atol=1e-6) def test_two_overlapping_binops_no_rerun(self): x = Tensor.randn(1,12,16,32) From e02da8f5acb496331f45a055ac694da41f3c4a3e Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 12:37:29 +0800 Subject: [PATCH 12/29] use op_in_parents (#12385) --- tinygrad/schedule/rangeify.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 08ca540a0f..38242849da 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -99,7 +99,7 @@ def extract_children(ctx:ChildrenContext, x:UOp): non_sink_children = [u for u in v if u.op is not Ops.SINK] if len(non_sink_children) <= 1: continue # NOTE: this gate shouldn't be here - if any(x.op is Ops.REDUCE_AXIS for x in k.toposort()) and any(x.op in {Ops.BUFFER, Ops.CONTIGUOUS} for x in k.toposort()): + if k.op_in_parents(Ops.REDUCE_AXIS) and k.op_in_parents(Ops.BUFFER, Ops.CONTIGUOUS): ctx.children[k] = non_sink_children def mark_children(ctx:ChildrenContext, x:UOp): From 57ad46c6e4013872031a2acf0562914c59943446 Mon Sep 17 00:00:00 2001 From: b1tg <33436708+b1tg@users.noreply.github.com> Date: Wed, 1 Oct 2025 12:56:45 +0800 Subject: [PATCH 13/29] rangeify: increase atol for test_two_binops_no_rerun passing on real windows machine (#12389) CPU_LLVM=1 --- test/unit/test_conv.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/unit/test_conv.py b/test/unit/test_conv.py index 47a24fa7f7..b9236798c2 100644 --- a/test/unit/test_conv.py +++ b/test/unit/test_conv.py @@ -51,8 +51,8 @@ class TestConv(unittest.TestCase): w = Tensor.randn(32,12,3,3) out = x.conv2d(w, stride=(2,2), padding=(1,1)) r1, r2 = out.relu(), (out-1) - np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-6) - np.testing.assert_allclose(r2.numpy(), out.numpy() - 1, atol=1e-6) + np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-5) + np.testing.assert_allclose(r2.numpy(), out.numpy() - 1, atol=1e-5) def test_two_overlapping_binops_no_rerun(self): x = Tensor.randn(1,12,16,32) From 714500edfdaf354d3c7fd71d2facfc2b5cd5d10d Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 1 Oct 2025 08:08:47 +0300 Subject: [PATCH 14/29] viz: add font-weight to OffscreenCanvas config (#12390) --- tinygrad/viz/js/worker.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/viz/js/worker.js b/tinygrad/viz/js/worker.js index 9fbcaa06e6..8c3703908b 100644 --- a/tinygrad/viz/js/worker.js +++ b/tinygrad/viz/js/worker.js @@ -2,7 +2,7 @@ const NODE_PADDING = 10; const LINE_HEIGHT = 14; const canvas = new OffscreenCanvas(0, 0); const ctx = canvas.getContext("2d"); -ctx.font = `${LINE_HEIGHT}px sans-serif`; +ctx.font = `350 ${LINE_HEIGHT}px sans-serif`; onmessage = (e) => { const { graph, additions } = e.data; From 05e91a248d2fbb5b842f1f8b946ed72e0787d88d Mon Sep 17 00:00:00 2001 From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com> Date: Wed, 1 Oct 2025 07:14:26 +0200 Subject: [PATCH 15/29] load alt value with cast (#12384) * add or_casted * add tests and fix old tests * cast load * move that to pm_render --- test/test_uop_graph.py | 15 +++++++++++++-- tinygrad/codegen/late/devectorizer.py | 5 +++++ tinygrad/uop/symbolic.py | 4 ---- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 2f6bc8d1a7..9f6d5c5ccb 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -420,7 +420,7 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_gated_load_fold(self): ridx0 = UOp.range(100, 0) d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) - ld = d0.index(ridx0, ridx0<50).load() + ld = d0.index(ridx0.valid(ridx0<50)).load() w = (ridx0<50).where(ld, 5) uops = to_uops_list([w]) for u in uops: @@ -430,13 +430,24 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_gated_load_folds_swapped_branches(self): ridx0 = UOp.range(100, 0) d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) - ld = d0.index(ridx0, (ridx0<50).logical_not()).load() + ld = d0.index(ridx0.valid((ridx0<50).logical_not())).load() w = (ridx0<50).where(5, ld) uops = to_uops_list([w]) for u in uops: assert u.op is not Ops.WHERE if u.op is Ops.LOAD: assert u.src[1].arg==5 + def test_where_on_gated_load_with_cast(self): + ridx0 = UOp.range(100, 0) + d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0) + gate_idx = ridx0.valid((ridx0<50)) + ld = d0.index(gate_idx).load().cast(dtypes.float) + w = (ridx0<50).where(ld, 5.0) + uops = to_uops_list([w]) + for u in uops: + assert u.op is not Ops.WHERE + if u.op is Ops.LOAD: assert u.src[1].arg == 5 + def test_where_in_store_becomes_gate(self): ridx0 = UOp.range(100, 0) d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index ade99dc849..d348319b3b 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -258,6 +258,11 @@ pm_render = PatternMatcher([ (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"), lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None), + # Where after gated load becomes alt value + (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(), + UPat.var("a")), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[1:]).cast(a.dtype)), + (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),), + allow_any_len=True, name="l").or_casted()), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[1:]).cast(a.dtype)), # gate any stores that aren't gated with ifs (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True), lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:]) if \ diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 90ebf5afd6..7a335339cf 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -505,10 +505,6 @@ sym = symbolic_flat+PatternMatcher([ # fold gated LOAD/STORE (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"), lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0 - (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l"), UPat.var("a")), - lambda c,idx,l,a: l.replace(src=(l.src[0], a)+l.src[1:])), - (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),), - allow_any_len=True, name="l")), lambda c,idx,l,a: l.replace(src=(l.src[0], a)+l.src[1:])), # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels (UPat(Ops.BARRIER, name="root"), lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg) From 42748ccb922c7ea073e4473fae20d2b7f984f44c Mon Sep 17 00:00:00 2001 From: b1tg <33436708+b1tg@users.noreply.github.com> Date: Wed, 1 Oct 2025 14:33:47 +0800 Subject: [PATCH 16/29] rangeify: fix test_prequant_conv2d_1x1 (#12391) --- test/test_quantize_onnx.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/test/test_quantize_onnx.py b/test/test_quantize_onnx.py index b15cdf12ce..005d978902 100644 --- a/test/test_quantize_onnx.py +++ b/test/test_quantize_onnx.py @@ -3,7 +3,6 @@ import numpy as np import unittest from dataclasses import replace from tinygrad import Tensor, Context, Device, dtypes -from tinygrad.helpers import RANGEIFY from tinygrad.uop.ops import Ops from tinygrad.codegen.opt import Opt, OptOps from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item, get_program @@ -94,8 +93,7 @@ class TestQuantizeOnnx(unittest.TestCase): X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8)) W = Tensor(np.random.uniform(0, 255, size=(64, 32, 1, 1)).astype(np.uint8)) out = X.conv2d(W, dtype=X.dtype) - # rangeify merges axis in a different order - opts = [Opt(op=OptOps.UPCAST, axis=0 if RANGEIFY else 1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] + opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] sexec(out, opts) def test_prequant_gemm(self): From 90b1c0dd9676e91376f577e8a65eeb7338f2eb12 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 1 Oct 2025 09:35:12 +0300 Subject: [PATCH 17/29] rangeify: test_where_fold kernel count (#12379) * rangeify: test_where_fold kernel count * get these from the index * replace ranges * fine * movement ops * diff * better --- test/test_assign.py | 3 +-- test/test_schedule.py | 9 +++++---- tinygrad/schedule/rangeify.py | 12 +++++++++--- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/test/test_assign.py b/test/test_assign.py index b35223fc12..63b6227b9e 100644 --- a/test/test_assign.py +++ b/test/test_assign.py @@ -379,8 +379,7 @@ class TestAssign(unittest.TestCase): a.assign(a + b) kc = GlobalCounters.kernel_count a.realize() - # rangeify makes two kernels - assert GlobalCounters.kernel_count - kc == (2 if RANGEIFY else 1) + assert GlobalCounters.kernel_count - kc == 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): diff --git a/test/test_schedule.py b/test/test_schedule.py index a0909a105f..d463c8aafb 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -1898,17 +1898,18 @@ class TestSchedule(unittest.TestCase): # NOTE: this is a bug on non rangeify np.testing.assert_equal(tst.numpy(), a.numpy()) - def test_setitem_sched(self, transpose=False): + def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1): a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize() - a2 = a.T if transpose else a + a2 = mop(a) expected = (a+a2).tolist() a.assign(a+a2) kcount = len(sched:=a.schedule()) run_schedule(sched) self.assertListEqual(a.tolist(), expected) - self.assertEqual(kcount, 2 if transpose else 1) + self.assertEqual(kcount, expected_kcount) @unittest.skipUnless(RANGEIFY>0, "this asserts on non rangeify") - def test_setitem_permuted_sched(self): self.test_setitem_sched(transpose=True) + 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) def test_sparse_categorical_crossentropy_simple(self): X = Tensor([[0, 2, 3], [1, 2, 3]]).realize() diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 38242849da..9af7b276ff 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -16,6 +16,13 @@ ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.ASSIGN, Ops.COPY, Ops.BUFFER, Ops.CONST, Ops.BIND, Ops.DEVICE, Ops.MSELECT, Ops.MSTACK, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.LOAD, Ops.KERNEL} +def find_permutes(a:UOp, b:UOp, assign:UOp): + if not (permutes:=[s for s in b.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS) + if s.op in GroupOp.Movement and s.op not in {Ops.RESHAPE, Ops.EXPAND, Ops.PAD, Ops.SHRINK}]): return + target = a.base + for p in permutes: + if any(s is target for s in p.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS-{Ops.BUFFER})): return assign.replace(src=(a, b.contiguous())) + earliest_rewrites = PatternMatcher([ # just removing it works... (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]), @@ -46,9 +53,8 @@ earliest_rewrites = PatternMatcher([ lambda x,target,assign: x.f(Ops.CONTIGUOUS, tag=assign.tag) if ((t:=target.base).op is not Ops.BUFFER and \ not (t.op is Ops.MSTACK and all(s.op is Ops.BUFFER for s in t.src))) else None), - # realize before assign if input permutes the target buffer - (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), lambda a,b,assign: assign.replace(src=(a, b.contiguous())) \ - if any(x.base is a.base and x is not a for x in b.toposort(gate=lambda x:x.op not in ALWAYS_CONTIGUOUS)) else None), + # realize before assign if input permutes the target buffer + (UPat(Ops.ASSIGN, src=(UPat.var("a"), UPat.var("b")), name="assign"), find_permutes), # copy only to different device (UPat(Ops.COPY, src=(UPat.var("x"), UPat()), name="copy"), lambda x,copy: x.f(Ops.NOOP, tag=copy.tag) if x.device == copy.device else None), From f205352cd756a782e9594476366f3748326c566b Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 14:43:29 +0800 Subject: [PATCH 18/29] remove ranges with 1s (#12388) * use op_in_parents * remove the ranges of 1 * fix CL image thing * fix realize --- tinygrad/schedule/rangeify.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 9af7b276ff..01eac77fcb 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -74,12 +74,14 @@ def realize_parents(ctx:dict[UOp, None], rb:UOp) -> None: def realize_assign(ctx:dict[UOp, None], a:UOp) -> None: if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None + # if it's a kernel, we don't realize it + if a.src[1].op is not Ops.KERNEL: ctx[a] = None do_realize = PatternMatcher([ # always realize SINK parents (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)), # always realize ASSIGN/COPY/BUFFER_VIEW/CONTIGUOUS - (UPat({Ops.ASSIGN, Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize), + (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize), # realize parents of COPY, MSELECT, MSTACK (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_parents), # realize input to assign (might be optimized out) @@ -131,7 +133,8 @@ class RangeifyContext: # create ranges range_idx: Iterator[int] = field(default_factory=itertools.count) - def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP): return UOp.range(s, next(self.range_idx), axistype) + def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP): + return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0) def map_reshape(idx:UOp, r:UOp): acc = 1 @@ -366,14 +369,15 @@ def cleanup_dead_axes(b:UOp): for s,rng in zip(b.shape, b.src[1:]): # skip for symbolic. TODO: fix this if rng.op is Ops.RANGE and rng.src[0].op is not Ops.CONST: return None - if rng not in b.src[0].sparents and rng.op is Ops.RANGE: + # CONSTs are already dead axes + if rng.op is Ops.CONST or (rng.op is Ops.RANGE and rng not in b.src[0].sparents): reshape.append(1) hit = True else: reshape.append(s) new_rng.append(rng) if hit: - # move the tag to the expand + # move the tag to the expand. NOTE: this expand tag might not survive return b.replace(src=b.src[0:1]+tuple(new_rng), tag=None).reshape(tuple(reshape)).expand(b.shape).replace(tag=b.tag) # if a buffer is being stored just for permutes or something, remove it @@ -381,7 +385,7 @@ def cleanup_dead_axes(b:UOp): def remove_bufferize(src:UOp, buf:UOp, idx:UOp): # see if we can't do it, should this ever hit? assert len(buf.src) == len(idx.src), "index on wrong bufferize" - assert all(x.op is Ops.RANGE for x in buf.src[1:]) + assert all(x.op in {Ops.RANGE, Ops.CONST} for x in buf.src[1:]) # if it's user contiguous, we never remove it if src.op in ALWAYS_RUN_OPS: return None @@ -411,7 +415,8 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp): # if it makes it here, the bufferize is removed # this is the ranges replaced - return src.substitute(dict(zip(buf.src[1:], idx.src[1:]))) + # NOTE: if buf src is a const, we don't replace it + return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST}) def pre_bufferize(b:UOp, x:UOp, copy:UOp): nb = b.replace(src=(b.src[0].contiguous(),)+b.src[1:]) @@ -519,8 +524,8 @@ def bufferize_to_store(x:UOp): ret = buf.reshape(shape).index(*rngs, dtype=sdtype).store(x.src[0], *rngs, dtype=x.dtype) ret = ret.forced_reshape(shape) # TODO: is this right? what if it's offset - if any(r.src[0].op is not Ops.RANGE for r in rngs): - sym_shape = tuple([ssimplify(r.src[0]) for r in rngs]) + if any(r.op is Ops.RANGE and r.src[0].op is not Ops.CONST for r in rngs): + sym_shape = tuple([ssimplify(r.src[0]) if r.op is not Ops.CONST else 1 for r in rngs]) ret = ret.shrink(tuple([(0,x) for x in sym_shape])) return ret.replace(tag=x.tag) From fe96c8d345013825899e9c986666c78cd5dd3fe3 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 1 Oct 2025 14:43:50 +0800 Subject: [PATCH 19/29] add HALF flag to tinygrad.apps.llm --- tinygrad/apps/llm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index 50649ffe4e..e331756eea 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -146,7 +146,7 @@ class Transformer: kv, state_dict = nn.state.gguf_load(gguf.to(None)) # all state items should be float16, not float32 - state_dict = {k:v.cast('float16') for k,v in state_dict.items()} + state_dict = {k:v.cast('float16') if getenv("HALF", 1) else v for k,v in state_dict.items()} # some models like Llama 3.2 don't have an output.weight, they just tie to the token_embd.weight if 'output.weight' not in state_dict: state_dict['output.weight'] = state_dict['token_embd.weight'] From 154d1143641bb326a8255278a487d3c2fdd0c049 Mon Sep 17 00:00:00 2001 From: b1tg <33436708+b1tg@users.noreply.github.com> Date: Wed, 1 Oct 2025 14:58:56 +0800 Subject: [PATCH 20/29] rangeify: fix abstractions2.py (#12386) * rangeify: fix abstractions2.py * tests * lint * only abstractions2 * base --- .github/workflows/test.yml | 3 +++ docs/abstractions2.py | 6 ++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 5b7cc33fa6..9f67dbac20 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -540,6 +540,9 @@ jobs: - name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1 run: | CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py + - name: Test Docs RANGEIFY=1 + run: | + RANGEIFY=1 python docs/abstractions2.py - name: Test const folding run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding" # RANGEIFY=2 isn't supported diff --git a/docs/abstractions2.py b/docs/abstractions2.py index e28037a768..cc23b27f6a 100644 --- a/docs/abstractions2.py +++ b/docs/abstractions2.py @@ -80,7 +80,9 @@ print("******** third, the UOp ***********") from tinygrad.engine.realize import run_schedule from tinygrad.engine.schedule import create_schedule_with_vars +from tinygrad.helpers import RANGEIFY from tinygrad.schedule.kernelize import get_kernelize_map +from tinygrad.schedule.rangeify import get_rangeify_map # allocate some values + load in values a = UOp.new_buffer(DEVICE, 1, dtypes.int32) @@ -93,10 +95,10 @@ out = a + b s = UOp(Ops.SINK, dtypes.void, (out,)) # group the computation into kernels -becomes_map = get_kernelize_map(s) +becomes_map = get_rangeify_map(s) if RANGEIFY else get_kernelize_map(s) # the compute maps to an assign -assign = becomes_map[a+b] +assign = becomes_map[a+b].base # the first source is the output buffer (data) assert assign.src[0].op is Ops.BUFFER From adc8c3b28f6ce66461c46f3dfb40c5ae29549f5e Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 1 Oct 2025 15:20:04 +0800 Subject: [PATCH 21/29] Revert "load alt value with cast (#12384)" (#12392) This reverts commit 05e91a248d2fbb5b842f1f8b946ed72e0787d88d. --- test/test_uop_graph.py | 15 ++------------- tinygrad/codegen/late/devectorizer.py | 5 ----- tinygrad/uop/symbolic.py | 4 ++++ 3 files changed, 6 insertions(+), 18 deletions(-) diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 9f6d5c5ccb..2f6bc8d1a7 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -420,7 +420,7 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_gated_load_fold(self): ridx0 = UOp.range(100, 0) d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) - ld = d0.index(ridx0.valid(ridx0<50)).load() + ld = d0.index(ridx0, ridx0<50).load() w = (ridx0<50).where(ld, 5) uops = to_uops_list([w]) for u in uops: @@ -430,24 +430,13 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_gated_load_folds_swapped_branches(self): ridx0 = UOp.range(100, 0) d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) - ld = d0.index(ridx0.valid((ridx0<50).logical_not())).load() + ld = d0.index(ridx0, (ridx0<50).logical_not()).load() w = (ridx0<50).where(5, ld) uops = to_uops_list([w]) for u in uops: assert u.op is not Ops.WHERE if u.op is Ops.LOAD: assert u.src[1].arg==5 - def test_where_on_gated_load_with_cast(self): - ridx0 = UOp.range(100, 0) - d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), (), 0) - gate_idx = ridx0.valid((ridx0<50)) - ld = d0.index(gate_idx).load().cast(dtypes.float) - w = (ridx0<50).where(ld, 5.0) - uops = to_uops_list([w]) - for u in uops: - assert u.op is not Ops.WHERE - if u.op is Ops.LOAD: assert u.src[1].arg == 5 - def test_where_in_store_becomes_gate(self): ridx0 = UOp.range(100, 0) d0 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(), (), 0) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index d348319b3b..ade99dc849 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -258,11 +258,6 @@ pm_render = PatternMatcher([ (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"), lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None), - # Where after gated load becomes alt value - (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(), - UPat.var("a")), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[1:]).cast(a.dtype)), - (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),), - allow_any_len=True, name="l").or_casted()), lambda c,idx,l,a: l.replace(src=(l.src[0], a.cast(l.dtype))+l.src[1:]).cast(a.dtype)), # gate any stores that aren't gated with ifs (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True), lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:]) if \ diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 7a335339cf..90ebf5afd6 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -505,6 +505,10 @@ sym = symbolic_flat+PatternMatcher([ # fold gated LOAD/STORE (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"), lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0 + (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l"), UPat.var("a")), + lambda c,idx,l,a: l.replace(src=(l.src[0], a)+l.src[1:])), + (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),), + allow_any_len=True, name="l")), lambda c,idx,l,a: l.replace(src=(l.src[0], a)+l.src[1:])), # remove VECTORIZE from SINK/BARRIER. TODO: SINK/BARRIER are really the same thing at GLOBAL/LOCAL levels (UPat(Ops.BARRIER, name="root"), lambda root: UOp(Ops.BARRIER, root.dtype, tuple(flatten(x.src if x.op in REMOVE_FROM_BARRIER else (x,) for x in root.src)), root.arg) From 689ab9151b78b29c9ef06b4e9efc364559fe16d5 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 1 Oct 2025 15:43:58 +0800 Subject: [PATCH 22/29] more RANGEIFY tests (#12393) would have caught the load alt regression without adding too many tests --- .github/workflows/test.yml | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9f67dbac20..4275067fe2 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -160,8 +160,10 @@ jobs: with: key: be-minimal deps: testing_minimal - - name: Test dtype with Python emulator - run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py + - name: Test dtype with Python emulator (with RANGEIFY) + run: | + RANGEIFY=0 DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py + RANGEIFY=1 DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py - name: Test ops with Python emulator run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20 - name: Test uops with Python emulator @@ -1037,3 +1039,9 @@ jobs: run: | python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT" python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20 + - name: Run pytest (${{ matrix.backend }}) with RANGEIFY + if: matrix.backend=='webgpu' + env: + RANGEIFY: 1 + shell: bash + run: python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20 From 6ba8bf282f426cc148e37aa02f7f16db9aadc246 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 1 Oct 2025 16:13:31 +0800 Subject: [PATCH 23/29] skip test_masked_select for RANGEIFY PYTHON (#12395) --- test/test_ops.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_ops.py b/test/test_ops.py index ddefcc1812..1609bd64ec 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -3164,7 +3164,8 @@ class TestOps(unittest.TestCase): helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf)) helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf)) - @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "AMD" and RANGEIFY, "very slow on MOCKGPU because reduce does not fold") + @unittest.skipIf(RANGEIFY and ((getenv("MOCKGPU") and Device.DEFAULT == "AMD") or Device.DEFAULT == "PYTHON"), + "very slow on MOCKGPU because reduce does not fold") def test_masked_select(self): helper_test_op([(32, 10)], lambda x: x.masked_select(x>0.5), lambda x: x.masked_select(x>0.5), forward_only=True) helper_test_op([(32, 10)], lambda x: x.masked_select(torch.tensor(True)), lambda x: x.masked_select(Tensor(True)), forward_only=True) From 6c95b1f39d126eb961ad4ed031456fa3bd469b47 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 1 Oct 2025 17:16:54 +0800 Subject: [PATCH 24/29] explicitly set device for CI unit test (#12399) --- .github/workflows/test.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4275067fe2..692b3a746c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -263,8 +263,10 @@ jobs: key: unittest-12 pydeps: "pillow numpy ftfy regex" deps: testing_unit + - name: Check Device.DEFAULT + run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT" - name: Run unit tests - run: python -m pytest -n=auto test/unit/ --durations=20 + run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20 - name: Run targetted tests on NULL backend run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py - name: Run SDXL on NULL backend From 60e52fbe3614902ada42e969439b6c6558ab424b Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 17:20:04 +0800 Subject: [PATCH 25/29] support opts in contig, simpler (#12400) --- test/test_opts.py | 3 +-- tinygrad/codegen/opt/postrange.py | 7 +++++-- tinygrad/schedule/rangeify.py | 15 ++++++++++++--- tinygrad/viz/serve.py | 4 ++-- 4 files changed, 20 insertions(+), 9 deletions(-) diff --git a/test/test_opts.py b/test/test_opts.py index 4a6310ef32..7bbdcfb61b 100644 --- a/test/test_opts.py +++ b/test/test_opts.py @@ -1,10 +1,9 @@ import unittest from tinygrad import Tensor, Device -from tinygrad.helpers import RANGEIFY, CPU_LLVM +from tinygrad.helpers import CPU_LLVM from tinygrad.codegen.opt import Opt, OptOps from tinygrad.engine.realize import get_program -@unittest.skipIf(RANGEIFY>0, "arg is partial contig in rangeify") class TestOpts(unittest.TestCase): def test_opt_upcast(self): opts = (Opt(OptOps.UPCAST, 0, 4),) diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 1fe4897242..10a415ec4f 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -5,7 +5,7 @@ from typing import cast, Final from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, can_pad, GroupOp from tinygrad.device import Buffer from tinygrad.dtype import AddrSpace, dtypes, ImageDType -from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts +from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters from tinygrad.codegen.simplify import pm_flatten_range from tinygrad.renderer import Renderer @@ -242,6 +242,9 @@ class Scheduler: if not (axis < len(axis_choices)): continue axes = list(axis_choices[axis]) + # tag the reduceop + self.ast = self.ast.substitute({reduceop: reduceop.replace(tag="TC")}) + # do optimizations and save the ranges try: for i,a in enumerate(axes): @@ -271,7 +274,7 @@ class Scheduler: if use_tensor_cores != 2: # fix the srcs - reduceop = [x for x in self.ast.toposort() if x.op is Ops.REDUCE][0] + reduceop = get_single_element([x for x in self.ast.toposort() if x.op is Ops.REDUCE and x.tag == "TC"]) tne = [x.replace(tag=1) for x in ne] ret = reduceop.substitute(dict(zip(ne, tne))) srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 01eac77fcb..06a930e64c 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -2,12 +2,13 @@ from typing import Any, cast, Iterator import functools, operator, itertools from dataclasses import dataclass, field from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace -from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify +from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, KernelInfo from tinygrad.uop.symbolic import sym, symbolic_simple from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup from tinygrad.schedule.kernelize import Kernel from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType from tinygrad.codegen.simplify import pm_flatten_range +from tinygrad.codegen.opt import Opt # ***************** # 0. do some cleanup rewrites, mostly copied from the old stuff @@ -555,6 +556,7 @@ class LocalAddBufferContext: vars:dict = field(default_factory=dict) range:int = 0 parent_tags:list = field(default_factory=list) + opts:tuple|None = None def debuf(ctx:LocalAddBufferContext, buf:UOp): ret = UOp(Ops.DEFINE_GLOBAL, buf.dtype.ptr(buf.arg), arg=ctx.dg) @@ -596,10 +598,16 @@ to_define_global = PatternMatcher([ (UPat(Ops.RANGE, name="r"), renumber_range), ]) +def get_contiguous(ctx:LocalAddBufferContext, x:UOp): + if isinstance(x.arg, tuple) and all(isinstance(y, Opt) for y in x.arg): ctx.opts = x.arg + return x.src[0] + rangeify_codegen = PatternMatcher([ + (UPat(Ops.CONTIGUOUS, name="x"), get_contiguous), + # no NOOP in the kernel graph # TODO: this can be moved into codegen? - (UPat((Ops.NOOP, Ops.CONTIGUOUS), name="x"), lambda x: x.src[0]), + (UPat(Ops.NOOP, name="x"), lambda x: x.src[0]), # strip the arg from store (UPat(Ops.STORE, name="x"), lambda x: x.replace(arg=None) if x.arg is not None else None), @@ -640,7 +648,8 @@ def split_store(ctx:list[UOp], x:UOp): metadatas = [ctx[y].metadata for y in lctx.parent_tags] # NOTE: the hack for COPY is here - ret = ret.sink() if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} else ret.src[1] + ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts) if lctx.opts is not None else None) \ + if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} else ret.src[1] kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1]) kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg) return x.as_buf().assign(kernel) diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index b3106c0241..7ac23c89f8 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -79,10 +79,10 @@ def uop_to_json(x:UOp) -> dict[int, dict]: arg = f"{x.arg:g}" if x.op is Ops.CONST and dtypes.is_float(x.dtype) else f"{x.arg}" label += f"\n{x.op.name}{idx} {arg}" + (f" {x.src[0].op}" if len(x.src) else "") try: + if len(rngs:=u.ranges): + label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})" if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: label += f"\n{shape_to_str(u.shape)}" - elif len(rngs:=u.ranges): - label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})" if u.op is Ops.INDEX: label += f"\n{u.render()}" except Exception: From ac3d457d5ead41a754dfd36178a271f0991dfae7 Mon Sep 17 00:00:00 2001 From: b1tg <33436708+b1tg@users.noreply.github.com> Date: Wed, 1 Oct 2025 17:58:19 +0800 Subject: [PATCH 26/29] rangeify: TestReduceOpsConstFolding (#12397) Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com> --- .github/workflows/test.yml | 2 +- tinygrad/codegen/simplify.py | 5 ++++- tinygrad/schedule/rangeify.py | 4 ++-- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 692b3a746c..98446abb3c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -548,7 +548,7 @@ jobs: run: | RANGEIFY=1 python docs/abstractions2.py - name: Test const folding - run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding" + run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded" # RANGEIFY=2 isn't supported #- name: Test CPU=1 RANGEIFY=2 # run: CPU=1 CPU_LLVM=0 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20 diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index b3d295287f..3d544f1f89 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -115,9 +115,12 @@ def reduce_unparented(red:UOp): for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count) return ret -pm_reduce_simplify = PatternMatcher([ +pm_reduce_unparented = PatternMatcher([ # remove any ranges from a REDUCE that aren't referenced in the reduce source (UPat(Ops.REDUCE, name="red"), reduce_unparented), +]) + +pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([ # remove REDUCE without loads (generic arange opt / indexing). TODO: support multi range (UPat(Ops.REDUCE, src=(UPat(), UPat()), name="red"), reduce_collapse), ]) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 06a930e64c..ab1c364cba 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -7,7 +7,7 @@ from tinygrad.uop.symbolic import sym, symbolic_simple from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup from tinygrad.schedule.kernelize import Kernel from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType -from tinygrad.codegen.simplify import pm_flatten_range +from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented from tinygrad.codegen.opt import Opt # ***************** @@ -699,7 +699,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]: # rangeify tsink = graph_rewrite(tsink, pm_rangeify, ctx=(rangeify_ctx:=RangeifyContext()), bottom_up=True, name="rangeify") # NOTE: sym (vs symbolic_simple) breaks things here because ranges with len 1 aren't handled right - tsink = graph_rewrite(tsink, symbolic_simple, name="symbolic") # this supports const folding + tsink = graph_rewrite(tsink, symbolic_simple+pm_reduce_unparented, name="symbolic") # this supports const folding tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers") tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rangeify_ctx, name="limit buffers") From f198a9e1ba4c504796e22d714835e43162243b95 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 1 Oct 2025 13:09:15 +0300 Subject: [PATCH 27/29] skip test_multihost_aware_schedule, assign devices mismatch (#12396) * minimal failing remote test * this should've never worked? * skip that test --- test/test_remote.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_remote.py b/test/test_remote.py index 796f98e62d..fd74dca65e 100644 --- a/test/test_remote.py +++ b/test/test_remote.py @@ -33,10 +33,11 @@ class TestRemoteMultiHost(unittest.TestCase): assert len(do.captured._jit_cache) == 1 and isinstance(do.captured._jit_cache[0].prg, RemoteGraph), repr(do.captured) @Context(JIT_BATCH_SIZE=2**32) + @unittest.skip("assign target and input devices mismatch") def test_multihost_aware_schedule(self): @TinyJit def do(*ts:Tensor): - acc = Tensor.zeros(1, dtype=dtypes.float32) + acc = Tensor.zeros(1, dtype=dtypes.float32).contiguous().realize() for t in ts: acc += t.sum() return acc.realize() From 74ee3059489dcb4e9baf0d6b41d40fc470ac4059 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 18:23:37 +0800 Subject: [PATCH 28/29] some rangeify tests fixed (#12403) --- test/test_tensor_uop.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/test_tensor_uop.py b/test/test_tensor_uop.py index 19541ead38..26925013ce 100644 --- a/test/test_tensor_uop.py +++ b/test/test_tensor_uop.py @@ -79,6 +79,7 @@ class TestTensorUOp(unittest.TestCase): np.testing.assert_allclose(out.numpy(), a.numpy()+b.numpy()+2) # NOTE: contiguous on a buffer collapses + @unittest.skip("contiguous on a buffer no longer collapses") def test_contiguous_empty(self): empty = Tensor.empty(1).contiguous() sched = empty.schedule() @@ -92,7 +93,7 @@ class TestTensorUOp(unittest.TestCase): out.realize() self.assertEqual(out.tolist(), Tensor.zeros(4, 8).tolist()) -reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, src=(UPat(), UPat(Ops.REDUCE_AXIS))))) +reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, src=(UPat(), UPat((Ops.REDUCE_AXIS, Ops.REDUCE)))))) class TestReduceOp(unittest.TestCase): def test_no_split_reduce_kernel(self): a = Tensor.rand(4, 4).realize() From 89bed2871619e485ef2cce6f6773918cae6fcbad Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 1 Oct 2025 18:45:16 +0800 Subject: [PATCH 29/29] split reduceop (#12404) * some rangeify tests fixed * bring split reduceop to rangeify * fix tests --- test/test_schedule.py | 8 ++++---- test/test_tensor_uop.py | 2 +- tinygrad/schedule/rangeify.py | 23 ++++++++++++++++++++++- 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/test/test_schedule.py b/test/test_schedule.py index d463c8aafb..4cf16351c0 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -1625,14 +1625,14 @@ class TestSchedule(unittest.TestCase): out = x.argmax(1) run_schedule(check_schedule(out, 2)) - def test_conv2d(self): _test_conv2d(4 if RANGEIFY else 7) - def test_conv2d_fused(self): _test_conv2d(4 if RANGEIFY else 5, FUSE_CONV_BW=1) + def test_conv2d(self): _test_conv2d(5 if RANGEIFY else 7) + def test_conv2d_fused(self): _test_conv2d(5 if RANGEIFY else 5, FUSE_CONV_BW=1) @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong") - def test_conv2d_half(self): _test_conv2d(4 if RANGEIFY else 7, dtype=dtypes.half) + def test_conv2d_half(self): _test_conv2d(5 if RANGEIFY else 7, dtype=dtypes.half) @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Causes other tests to fail") - @unittest.expectedFailure + @unittest.skipIf(not RANGEIFY, "passes on RANGEIFY") def test_conv2d_fused_half(self): _test_conv2d(5, dtype=dtypes.half) def test_schedule_mem_used(self): diff --git a/test/test_tensor_uop.py b/test/test_tensor_uop.py index 26925013ce..72c9f3a661 100644 --- a/test/test_tensor_uop.py +++ b/test/test_tensor_uop.py @@ -93,7 +93,7 @@ class TestTensorUOp(unittest.TestCase): out.realize() self.assertEqual(out.tolist(), Tensor.zeros(4, 8).tolist()) -reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, src=(UPat(), UPat((Ops.REDUCE_AXIS, Ops.REDUCE)))))) +reduce_kernel = UPat(Ops.SINK, src=(UPat(Ops.STORE, allow_any_len=True, src=(UPat(), UPat((Ops.REDUCE_AXIS, Ops.REDUCE)))))) class TestReduceOp(unittest.TestCase): def test_no_split_reduce_kernel(self): a = Tensor.rand(4, 4).realize() diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index ab1c364cba..649a67a64b 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute, ssimplify, KernelInfo from tinygrad.uop.symbolic import sym, symbolic_simple -from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup +from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, RANGEIFY, Context, flatten, dedup, unwrap, all_int, DEBUG, SPLIT_REDUCEOP from tinygrad.schedule.kernelize import Kernel from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_unparented @@ -24,10 +24,31 @@ def find_permutes(a:UOp, b:UOp, assign:UOp): for p in permutes: if any(s is target for s in p.toposort(gate=lambda s:s.op not in ALWAYS_CONTIGUOUS-{Ops.BUFFER})): return assign.replace(src=(a, b.contiguous())) +def split_reduceop(reduce:UOp, x:UOp): + if prod(reduce.shape) == 0: return None + if not SPLIT_REDUCEOP or not all_int(x.shape) or (prod(x.shape)//prod(reduce.shape))= 3: print(f"split {divisor}: {x.shape} -> {splitted.shape} -> {reduce.shape}") + # reduce original axes, then split + return splitted.r(*reduce.arg).contiguous().r(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape).replace(tag=reduce.tag) + earliest_rewrites = PatternMatcher([ # just removing it works... (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD, Ops.FUSE), name="x"), lambda x: x.src[0]), + # split_reduceop + (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)), split_reduceop), + # preserve tags? # reduce of size 0 is the identity element (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),