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)))) diff --git a/test/test_outerworld.py b/test/test_outerworld.py index d65386f470..a3674f4613 100644 --- a/test/test_outerworld.py +++ b/test/test_outerworld.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Tensor, UOp, Variable -from tinygrad.uop.ops import Ops, AxisType +from tinygrad.uop.ops import AxisType, Ops #@unittest.skip("TODO: understand assign") class TestOuterworldAssign(unittest.TestCase): diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 867bea76bb..812fdfbc94 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -111,6 +111,8 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([ # endrange/reduce for outerworld range work (UPat(Ops.ENDRANGE, src=(UPat(Ops.RANGE),), allow_any_len=True), lambda: True), + + # 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:])), ]) diff --git a/tinygrad/viz/js/index.js b/tinygrad/viz/js/index.js index c0b461b560..92223b6682 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(); @@ -170,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(""); @@ -281,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 @@ -350,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 @@ -443,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 ?? ""); });