From 464c56862f9a23f0cfb58b59b8b77e9bd4e1c8a3 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 10 Oct 2025 13:58:58 +0300 Subject: [PATCH 1/4] viz: update ansi regex (#12605) * viz: update ansi regex * better * add ansi_colors_light * javascript --- tinygrad/viz/js/index.js | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index c0b461b560..11b491e4a9 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -14,8 +14,9 @@ const darkenHex = (h, p = 0) => ).toString(16).padStart(6, '0')}`; const ANSI_COLORS = ["#b3b3b3", "#ff6666", "#66b366", "#ffff66", "#6666ff", "#ff66ff", "#66ffff", "#ffffff"]; +const ANSI_COLORS_LIGHT = ["#d9d9d9","#ff9999","#99cc99","#ffff99","#9999ff","#ff99ff","#ccffff","#ffffff"]; const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g), - ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? ANSI_COLORS[(parseInt(code)-30+60)%60] : defaultColor })); + ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? (code>=90 ? ANSI_COLORS_LIGHT : ANSI_COLORS)[(parseInt(code)-30+60)%60] : defaultColor })); const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect(); From a62dc9ceb5b8f42ac551f864e98e2d8d5e2158da Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 10 Oct 2025 14:07:30 +0300 Subject: [PATCH 2/4] viz: light up buffer path (#12603) --- tinygrad/viz/js/index.js | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index 11b491e4a9..92223b6682 100644 --- a/tinygrad/viz/js/index.js +++ b/tinygrad/viz/js/index.js @@ -171,7 +171,7 @@ function tabulate(rows) { return root; } -var data, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity; +var data, focusedDevice, focusedShape, canvasZoom, zoomLevel = d3.zoomIdentity; async function renderProfiler() { displayGraph("profiler"); d3.select(".metadata").html(""); @@ -282,7 +282,7 @@ async function renderProfiler() { const html = document.createElement("div"); const rows = [["DType", dtype], ["Len", formatUnit(sz)], ["Size", formatUnit(nbytes, "B")], ["Lifetime", formatTime(dur)]]; const info = html.appendChild(tabulate(rows).node()); - const arg = {tooltipText:info.outerHTML, html}; + const arg = {tooltipText:info.outerHTML, html, key:`${k}-${num}`}; shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) }); } // generic polygon merger @@ -351,6 +351,7 @@ async function renderProfiler() { for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]); ctx.closePath(); ctx.fillStyle = e.fillColor; ctx.fill(); + if (focusedShape && e.arg?.key === focusedShape) { ctx.lineWidth = 1.4; ctx.strokeStyle = "#c9a8ff"; ctx.stroke(); } continue; } // contiguous rect @@ -444,6 +445,7 @@ async function renderProfiler() { e.preventDefault(); const foundRect = findRectAtPosition(e.clientX, e.clientY); if (foundRect?.step != null) return setCtxWithHistory(foundRect.ctx, foundRect.step); + if (foundRect?.key != focusedShape) { focusedShape = foundRect?.key; render(zoomLevel); } return document.querySelector(".metadata").replaceChildren(foundRect?.html ?? ""); }); From 001b3710d357210cfb260cce34a084ab3b8af1a6 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 10 Oct 2025 19:23:21 +0800 Subject: [PATCH 3/4] enable some test_ops tests (#12607) --- test/test_ops.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index bbc33147ee..6e413d3dff 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings import numpy as np from typing import List, Callable import torch -from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM +from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LLVM from tinygrad import Tensor, Device, dtypes from tinygrad.tensor import _to_np_dtype from tinygrad.device import is_dtype_supported @@ -901,7 +901,6 @@ class TestOps(unittest.TestCase): def test_abs_exact(self): helper_test_op(None, torch.abs, Tensor.abs, vals=[[-1.,0,1]]) - @unittest.skipIf(TRANSCENDENTAL and Device.DEFAULT=="AMD", "TODO: remu crashes") def test_log(self): helper_test_op([(45,65)], torch.log, Tensor.log) helper_test_op(None, torch.log, Tensor.log, vals=[[math.inf, -math.inf, math.nan]]) @@ -911,7 +910,6 @@ class TestOps(unittest.TestCase): helper_test_op(None, torch.log2, Tensor.log2, vals=[[math.inf, -math.inf, math.nan]]) helper_test_op([()], torch.log2, Tensor.log2) - @unittest.skipIf(TRANSCENDENTAL and Device.DEFAULT=="AMD", "TODO: remu crashes") def test_exp(self): helper_test_op([(45,65)], torch.exp, Tensor.exp) helper_test_op(None, torch.exp, Tensor.exp, vals=[[math.inf, -math.inf, math.nan]]) @@ -1549,7 +1547,6 @@ class TestOps(unittest.TestCase): helper_test_op([(3,4,5,6)], lambda x: torch.stack(torch.std_mean(x, axis=(1,2))), lambda x: Tensor.stack(*x.std_mean(axis=(1,2)))) - @unittest.skip("TODO: this fails because of loaded nan in mul folding") def test_std_mean_loaded_nan(self): helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))), lambda x: Tensor.stack(*x.std_mean(axis=(1,3)))) From 7596c1b8f5bc9fe646366a4e4921af194377f943 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 10 Oct 2025 20:06:41 +0800 Subject: [PATCH 4/4] TestOuterworldReduce works (#12608) --- test/test_outerworld.py | 61 +++++++++++++++++------------------ tinygrad/schedule/rangeify.py | 2 +- tinygrad/uop/spec.py | 3 ++ 3 files changed, 34 insertions(+), 32 deletions(-) diff --git a/test/test_outerworld.py b/test/test_outerworld.py index 449d122017..5714cc0748 100644 --- a/test/test_outerworld.py +++ b/test/test_outerworld.py @@ -1,5 +1,15 @@ import unittest -from tinygrad import Tensor, UOp, GlobalCounters, Context +from tinygrad import Tensor, UOp +from tinygrad.uop.ops import AxisType, Ops + +class TestOuterworldReduce(unittest.TestCase): + def test_reduce(self): + x = Tensor.ones(5, 5).contiguous() + a = UOp.range(5, -1, AxisType.REDUCE) + out = x[a] + # TODO: syntax for this + t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, a), arg=Ops.ADD)) + self.assertListEqual(t.tolist(), [5.,5.,5.,5.,5.]) class TestOuterworld(unittest.TestCase): def test_range_plus_1(self): @@ -13,6 +23,17 @@ class TestOuterworld(unittest.TestCase): self.assertTrue((t+1==cpy).all().item()) + def test_range_plus_1_transpose(self): + t = Tensor.arange(100).reshape(10,10).realize() + + # passthrough ranges + a = UOp.range(10, -1) + sel = t[a] + 1 + assert sel.shape == (10,) + cpy = sel.reshape(10, 1).expand(10, a).contiguous().realize() + + self.assertTrue(((t+1).T==cpy).all().item()) + def test_flip_range(self): t = Tensor.rand(10, 10).realize() @@ -37,39 +58,17 @@ class TestOuterworld(unittest.TestCase): out.realize() self.assertTrue((out==20).all().item()) - @unittest.skip("opts don't work") - def test_triple_gemm(self): - x = Tensor.rand(1, 16).realize() - W = Tensor.rand(3, 16, 16).realize() + def test_fancy_vmap(self): + def f(x,y): return x+y - manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize() + x = Tensor.arange(9).reshape(3,3).contiguous() + y = Tensor.arange(9).reshape(3,3).contiguous() a = UOp.range(3, -1) - x = x.assign(x @ W[a]) - out = x.contiguous(a)[-1].contiguous().realize() - - self.assertTrue((manual==out).all().item()) - - def test_setitem_pyrange(self): - with Context(DEBUG=0): - t = Tensor.rand(10).realize() - o = Tensor.empty(10) - GlobalCounters.reset() - for i in range(10): - o[i] = t[i] - o.realize() - self.assertTrue((t==o).all().item()) - - @unittest.skip("TODO: fix this") - def test_setitem(self): - with Context(DEBUG=0): - t = Tensor.rand(10).realize() - o = Tensor.empty(10) - GlobalCounters.reset() - i = UOp.range(10, -1) - o[i] = t[i] - o.contiguous(i).realize() - self.assertTrue((t==o).all().item()) + out = f(x[:,a], y[a,:]) + # TODO: this should support flatten + out = out.reshape(1, 3).expand(a, 3).contiguous().realize() + self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist()) if __name__ == '__main__': unittest.main() \ No newline at end of file diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index bffff16e70..5eb06fd7ee 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -442,7 +442,7 @@ def tag_uop(ctx:list[UOp], x:UOp): add_tags = PatternMatcher([ # don't tag BUFFERs, they are global (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND, - Ops.MSTACK, Ops.MSELECT}.union(GroupOp.Movement), name="x"), tag_uop), + Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop), (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)), ]) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 9bd74a52c9..aadacfb76e 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -108,6 +108,9 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([ (UPat(Ops.COPY, name="copy", src=(UPat.var("x"), UPat(Ops.DEVICE)), arg=None), lambda copy,x: copy.dtype == x.dtype), (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)), (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)), + + # REDUCE with an outerworld range + (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])), ]) # ***** uop type spec *****