From b5588d341b50c4f3fd363682e8c43eb17fc6a66a Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Wed, 18 Feb 2026 19:49:47 -0800 Subject: [PATCH 1/4] uop_given_valid fixes many gated reads for IMAGE=1 (#14877) * add replay script * pkl is arg * that needs uop_given_valid * cleanup --- tinygrad/codegen/late/devectorizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index dd1baa7e6a..daf92cbc44 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -189,7 +189,7 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]: h, w = dt.shape[0], dt.shape[1] if IMAGE == 1 and valid is not None and (tp:=dt.size // 4) // 64: h, w = max(([(1, tp)] * (tp < 16384)) + [(tp//64//k, 64*k) for k in range(ceildiv(tp//64, 16384), min(tp//64, 256)+1) if (tp//64) % k == 0], - key=lambda hw: len(_drop_valid_stmts(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1])), *hw))) + key=lambda hw: len(_drop_valid_stmts(valid, uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw))) buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4), w * 4 * dt.itemsize)) oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % w, (x // (4*w)))) return x, idx.replace(src=(buf, oidx.valid(valid))), w, h From 8d8da185ec85dc093afb31efe41e603798ff6b2c Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 18 Feb 2026 22:53:28 -0500 Subject: [PATCH 2/4] minor handle_allreduce cleanup [pr] (#14876) no more lbs, also use a divmod --- tinygrad/schedule/multi.py | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index e79c34e6c8..7d3b3fada3 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -8,36 +8,36 @@ from tinygrad.dtype import dtypes def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: if not isinstance(buf.device, tuple): return None assert all_int(buf.shape), f"does not support symbolic shape {buf.shape}" - n_lbs, shape, numel = len(buf.device), buf.shape, prod(buf.shape) + ndev, shape, numel = len(buf.device), buf.shape, prod(buf.shape) # ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically) # fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks. - use_all2all = (ALL2ALL >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1)) - use_ring = not use_all2all and (RING >= 2 or (n_lbs > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1)) - if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {n_lbs}x{numel} | {buf.dtype}") + use_all2all = (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1)) + use_ring = not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1)) + if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}") # contiguous before we copy it buf = buf.contiguous() # naive: copy to all devices. if you shrink later, that'll be handled if not use_ring and not use_all2all: - return functools.reduce(lambda x,y: x.alu(red.arg, y), [UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(n_lbs)]) + return functools.reduce(lambda x,y: x.alu(red.arg, y), [UOp(Ops.COPY, buf.dtype, (buf.mselect(i), red.src[1])) for i in range(ndev)]) - # chunk data into n_lbs pieces + # chunk data into ndev pieces factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1) - base, left = (numel // factor) // n_lbs, (numel // factor) % n_lbs - chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (n_lbs - left), initial=0))) + base, left = divmod(numel // factor, ndev) + chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (ndev - left), initial=0))) # reduce-scatter reduced_chunks = [] for i,(s,e) in enumerate(chunks): if use_all2all: - chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(n_lbs)] + chunks_on_i = [buf.mselect(j).reshape((numel,)).shrink(((s,e),)).copy_to_device(buf.device[i]) for j in range(ndev)] reduced_chunks.append(functools.reduce(lambda x,y: x.alu(red.arg, y), chunks_on_i)) else: chunk, reduced = buf.reshape((numel,)).shrink(((s,e),)), buf.reshape((numel,)).shrink(((s,e),)) - for step in range(n_lbs-1): - src, dest = (i+step)%n_lbs, (i+step+1)%n_lbs + for step in range(ndev-1): + src, dest = (i+step)%ndev, (i+step+1)%ndev cp = reduced.copy_to_device(buf.device[dest], src if isinstance(reduced.device, tuple) else None) reduced = cp.alu(red.arg, chunk.copy_to_device(buf.device[dest], dest)) reduced_chunks.append(reduced) @@ -46,12 +46,12 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None: copied_chunks = [] for i,rc in enumerate(reduced_chunks): if isinstance(red.src[1].arg, str): copied_chunks.append(rc.copy_to_device(red.src[1].arg)) - elif use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(n_lbs)))) + elif use_all2all: copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(rc.copy_to_device(buf.device[j]) for j in range(ndev)))) else: - this_chunk: list[UOp|None] = [None] * n_lbs - this_chunk[(i+n_lbs-1)%n_lbs] = rc - for step in range(n_lbs-1): - this_chunk[(i+step)%n_lbs] = rc = rc.copy_to_device(buf.device[(i+step)%n_lbs]) + this_chunk: list[UOp|None] = [None] * ndev + this_chunk[(i+ndev-1)%ndev] = rc + for step in range(ndev-1): + this_chunk[(i+step)%ndev] = rc = rc.copy_to_device(buf.device[(i+step)%ndev]) copied_chunks.append(UOp(Ops.MSTACK, buf.dtype, tuple(cast(list[UOp], this_chunk)))) # reassemble From 5bc65ec669b06de9adbbe5abc3c89099f137818c Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Thu, 19 Feb 2026 12:08:26 +0800 Subject: [PATCH 3/4] applied_opts/estimates in program spec are aliases for the sink arg (#14860) * remove applied_opts from programspec * comment that out * placement * update tests * p.ast.arg * remove todo comment * maybe this too * it can exist as an alias, also for estimates --- tinygrad/renderer/__init__.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 28b9f800e7..1778efa4fe 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -80,7 +80,7 @@ class ProgramSpec: ins:list[int]=field(default_factory=list) @property - def estimates(self) -> Estimates: return self.ast.arg.estimates + def estimates(self) -> Estimates: return self.ast.arg.estimates if self.ast.arg is not None and self.ast.arg.estimates is not None else Estimates() @functools.cached_property def function_name(self) -> str: return to_function_name(self.name) @@ -89,10 +89,7 @@ class ProgramSpec: def runtimevars(self) -> dict[str, int]: return {v.arg[0]: i for i, v in enumerate(self.vars) if v.arg[0] == 'core_id'} @property - def applied_opts(self) -> tuple[Opt, ...]|None: - if self.uops is None: return None - assert self.uops[-1].op is Ops.SINK, self.uops[-1].op - return self.uops[-1].arg.applied_opts + def applied_opts(self) -> tuple[Opt, ...]|None: return self.ast.arg.applied_opts if self.ast.arg is not None else None def launch_dims(self, var_vals:dict[str, int]): global_size = [sym_infer(sz, var_vals) for sz in self.global_size] From 2f0f8b5776fc38184cc38c6e27b44166d9676625 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 19 Feb 2026 14:23:28 +0800 Subject: [PATCH 4/4] more test relaxations from prealloc_bufs (#14880) --- test/backend/test_schedule.py | 30 +++++++++++---------- test/null/test_opts.py | 21 --------------- test/null/test_schedule.py | 6 ++++- test/null/test_tensor_uop_representation.py | 1 + test/null/test_uops_stats.py | 1 + tinygrad/uop/ops.py | 2 ++ 6 files changed, 25 insertions(+), 36 deletions(-) delete mode 100644 test/null/test_opts.py diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index b7332275f8..b507642b29 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -168,13 +168,13 @@ class TestSchedule(unittest.TestCase): a = Tensor.full((4,), 4.0).contiguous().realize() b = Tensor.full((4,), 2.0).contiguous().realize() expr = (a*b)/b - check_schedule(expr, 0) + run_schedule(check_schedule(expr, 0)) np.testing.assert_allclose(expr.numpy(), np.full((4,), 4.0)) def test_div_collapse_const(self): a = Tensor.full((4,), 4.0).contiguous().realize() expr = a/a - check_schedule(expr, 0) + run_schedule(check_schedule(expr, 0)) np.testing.assert_allclose(expr.numpy(), np.full((4,), 1.0)) def test_div_collapse(self): @@ -1236,11 +1236,12 @@ class TestView(unittest.TestCase): bv = b.pad(((0, 2),))[-2:] # this becomes a late a*0 late_mul = a*bv - check_schedule(late_mul, 0) + run_schedule(check_schedule(late_mul, 0)) + # NOTE: no longer checked # the arange doesn't realize - self.assertIsNone(b.uop.base.realized) + #self.assertIsNone(b.uop.base.realized) # mul doesn't realize - self.assertIsNone(late_mul.uop.base.realized) + #self.assertIsNone(late_mul.uop.base.realized) self.assertEqual(late_mul.tolist(), [0, 0]) # SINK has two branches: @@ -1256,17 +1257,18 @@ class TestView(unittest.TestCase): s = check_schedule([late_mul, other_child], 2) # the arange becomes a BUFFER self.assertIs(b.uop.base.op, Ops.BUFFER) + # NOTE: no longer checked # mul still collapses - self.assertIs(late_mul.uop.base.op, Ops.CONST) + #self.assertIs(late_mul.uop.base.op, Ops.CONST) run_schedule(s) self.assertEqual(other_child.tolist(), [2, 3, 4]) @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu") class TestCopyFolding(unittest.TestCase): def test_const_copy_is_free(self): - b = Tensor(1).to("CPU") - check_schedule(b, 0, filter_sink=False) - assert b.item() == 1 + b = Tensor(1).to("CPU") * 4 + run_schedule(check_schedule(b, 0, filter_sink=False)) + assert b.item() == 4 def test_one_hot_with_copy(self): y = Tensor([1, 2, 3]).to("CPU") @@ -1274,16 +1276,16 @@ class TestCopyFolding(unittest.TestCase): check_schedule(x, 3, filter_sink=False) def test_const_copy_multi(self): - x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) - check_schedule(x, 0, filter_sink=False) - self.assertEqual(x.item(), 1) + x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) * 2 + run_schedule(check_schedule(x, 0, filter_sink=False)) + self.assertEqual(x.item(), 2.0) def test_late_const_copy_folding(self): a = Tensor.arange(3).realize() zeros = Tensor.zeros(3).realize() - b = (a*zeros).to("CPU") + b = (a*zeros).to("CPU") + 1 run_schedule(check_schedule(b, 0, filter_sink=False)) - self.assertListEqual(b.tolist(), [0, 0, 0]) + self.assertListEqual(b.tolist(), [1, 1, 1]) self.assertEqual(b.device, "CPU") def test_alu_after_copy(self): diff --git a/test/null/test_opts.py b/test/null/test_opts.py deleted file mode 100644 index 359441cbf1..0000000000 --- a/test/null/test_opts.py +++ /dev/null @@ -1,21 +0,0 @@ -import unittest -from tinygrad import Tensor, Device -from tinygrad.helpers import CPU_LLVM, CPU_LVP -from tinygrad.codegen.opt import Opt, OptOps -from tinygrad.engine.realize import get_program - -class TestOpts(unittest.TestCase): - def test_opt_upcast(self): - opts = (Opt(OptOps.UPCAST, 0, 4),) - a = Tensor.empty(16) - b = Tensor.empty(16) - out = (a+b).contiguous(arg=opts) - s = out.schedule() - self.assertEqual(s[-1].ast.arg.opts_to_apply, opts) - if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP: - prg = get_program(s[-1].ast, renderer=Device[Device.DEFAULT].renderer) - self.assertIn('float4', prg.src) - -if __name__ == '__main__': - unittest.main() - diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index 03785adfc2..edb5d470ee 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -74,7 +74,7 @@ class TestBufferUOp(unittest.TestCase): self.assertIsNotNone(a.uop.buffer) def test_const_does_not_realize(self): - a = Tensor(1)+Tensor(2) + a = Tensor(1) run_schedule(check_schedule(a, 0)) self.assertIsNone(a.uop.base.realized) @@ -642,6 +642,7 @@ class TestSchedule(unittest.TestCase): t = Tensor([1.0, 2.0, 3.0]) ** 8 self.assertEqual(self._alu_from_tensor(t), [Ops.MUL, Ops.MUL, Ops.MUL]) + @unittest.skip("const folding is removed") def test_pow_const_tensor_to_zero(self): x = Tensor([1,2,3,4]) out = x ** Tensor(0.0) @@ -1003,6 +1004,7 @@ class TestUOpBecome(unittest.TestCase): # sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer + @unittest.skip("no longer supported") def test_reorder_expand(self): a = Tensor.empty(4, 1) b = a.expand(4, 4).reciprocal() @@ -1038,6 +1040,7 @@ class TestUOpBecome(unittest.TestCase): late_add = noop+2 late_add.realize() + @unittest.skip("const folding is removed") def test_become_const_in_base(self): a = Tensor.empty(4) b = a*0 @@ -1045,6 +1048,7 @@ class TestUOpBecome(unittest.TestCase): check_schedule(b, 0) assert UPat(Ops.CONST, arg=0).match(b.uop.base, {}) # scheduling replaces the tensor uop with a VIEW(BUFFER) + @unittest.skip("const folding is removed") def test_become_const_from_const(self): const_add = Tensor(1)+Tensor(2) assert UPat(Ops.ADD).match(const_add.uop, {}) diff --git a/test/null/test_tensor_uop_representation.py b/test/null/test_tensor_uop_representation.py index e3b173d639..f393ecd2ff 100644 --- a/test/null/test_tensor_uop_representation.py +++ b/test/null/test_tensor_uop_representation.py @@ -8,6 +8,7 @@ def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pa def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat) class TestTensorMutates(unittest.TestCase): + @unittest.skip("this doesn't mutate anymore") def test_mutate_add(self): a = Tensor([1,2,3]) b = Tensor([4,5,6]) diff --git a/test/null/test_uops_stats.py b/test/null/test_uops_stats.py index 67de94e4d0..94d44759d6 100644 --- a/test/null/test_uops_stats.py +++ b/test/null/test_uops_stats.py @@ -45,6 +45,7 @@ class TestMemoryCount(unittest.TestCase): _, mem = get_stats(a+b) self.assertEqual(mem, 1024*1024*2 + 1024) # 1 full read + 1 lil read + 1 write + @unittest.skip("no longer supported") def test_both_expanded(self): # TODO: this probably should be a full write a = Tensor.empty(1024, 1, dtype=dtypes.uint8).expand(1024, 1024) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 8a7e5b60dc..b5e5809c96 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -506,6 +506,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass): @functools.cached_property def axis(self) -> int|None: + # COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK + if self.op is Ops.COPY: return None if self.op is Ops.MULTI: return self.arg # NOTE: they all have to share an axis, we always choose [-1] if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None