minimal new linearizer

* cleanup

* cleanup there

test/unit/test_helpers.py

@@ -1,9 +1,8 @@
 import ctypes, gzip, unittest, timeit
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address, get_contraction
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.tensor import Tensor, get_shape
-from tinygrad.shape.view import get_contraction
 import numpy as np
 
 VARIABLE = ContextVar("VARIABLE", 0)

test/unit/test_shapetracker.py

@@ -154,21 +154,6 @@ class TestRealStrides(unittest.TestCase):
     ))
     self.assertEqual(st.is_expanded(), (False, False, False, True, False))
 
-class TestRealSimplifies(unittest.TestCase):
-  def tearDown(self):
-    self.st = self.st.simplify()
-    assert len(self.st.views) == 1
-
-  def test_1(self):
-    self.st = ShapeTracker((
-      View.create((1, 3, 2, 11, 4, 28), (0, 308, 0, 28, 0, 1), 0, None),
-      View.create((1, 3, 2, 11, 26, 1, 1, 3), (0, 2464, 0, 112, 1, 0, 0, 29), 0, None)))
-
-  def test_2(self):
-    self.st = ShapeTracker((
-      View.create((8, 3, 3, 11, 2, 28), (924, 308, 0, 28, 0, 1), 0, None),
-      View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)))
-
 class TestIndexExpressions2d(unittest.TestCase):
   def setUp(self):
     shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5

test/unit/test_shapetracker_math.py

@@ -62,6 +62,7 @@ class TestShapeTrackerAdd(unittest.TestCase):
     b = ShapeTracker.from_shape((100,))
     assert a+b == b
 
+  @unittest.skip("no longer simplifies")
   def test_simple_add_permute(self):
     a = ShapeTracker.from_shape((10, 10))
     a = a.permute((1,0))

test/unit/test_symbolic_shapetracker.py

@@ -28,7 +28,7 @@ class TestSymbolic(unittest.TestCase):
   def test_merge_view_recursion_err(self):
     vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
     vm1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
-    self.assertEqual(vm2+vm1, vm1)
+    self.assertEqual(vm2+vm1, None)
 
   def test_merge_view_recursion_err2(self):
     vm2 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(0,), offset=0, mask=None, contiguous=False)

test/unit/test_view.py

@@ -69,161 +69,5 @@ class TestMergeDims(unittest.TestCase):
     # print(f"{ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))).views[-1]}")
     self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))
 
-class TestMergeViews(unittest.TestCase):
-  def test_with_mask_0(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    v0 = View(shape=(1, 1, 5, 8), strides=(0, 0, 5, 1), offset=-3, mask=((0, 1), (0, 1), (0, 5), (3, 8)), contiguous=False)
-    v1 = View(shape=(1, 1, 2, 2), strides=(0, 0, 8, 1), offset=3, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 1, 2, 2), strides=(0, 0, 5, 1), offset=0, mask=None, contiguous=False))
-
-  def test_with_mask_1(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, 3, 1), offset=-6, mask=((0, 3), (0, 3), (2, 4), (1, 3)), contiguous=False)
-    v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, 3, 1), offset=7, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=1, mask=None, contiguous=False))
-
-  def test_with_mask_2(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (0, 2)), contiguous=False)
-    v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, -3, 1), offset=3, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=3, mask=None, contiguous=False))
-
-  def test_with_mask_3(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    # has a mask in the final view
-    v0 = View(shape=(3, 3, 4, 4), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False)
-    v1 = View(shape=(3, 3, 4, 2), strides=(48, 16, 4, 1), offset=0, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 4, 2), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False))
-
-  def test_with_mask_4(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    # has a mask in the final view
-    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (1, 3)), contiguous=False)
-    v1 = View(shape=(3, 3, 3, 3), strides=(45, 15, 3, 1), offset=6, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 3, 3), strides=(0, 0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0), (0, 0)), contiguous=False))
-
-  def test_with_mask_5(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    # has a mask in the final view
-    v0 = View(shape=(1, 1, 6, 5), strides=(0, 0, 5, 1), offset=-5, mask=((0, 1), (0, 1), (1, 6), (0, 5)), contiguous=False)
-    v1 = View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=3, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=-2, mask=((0, 1), (0, 1), (1, 6), (0, 3)), contiguous=False))
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer1(self):
-    v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False)
-    v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False)
-    target = View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer2(self):
-    v0 = View(shape=(5, 10, 12), strides=(100, 1, 10), offset=-20, mask=((0, 5), (0, 10), (2, 12)), contiguous=False)
-    v1 = View(shape=(10, 6, 5, 2, 2), strides=(12, 2, 120, 1, 0), offset=0, mask=None, contiguous=False)
-    target = View(shape=(10, 6, 5, 2, 2), strides=(1, 20, 100, 10, 0), offset=-20, mask=((0, 10), (1, 6), (0, 5), (0, 2), (0, 2)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer3(self):
-    v0 = View(shape=(8, 7, 3), strides=(1, 12, -4), offset=6, mask=((2, 6), (0, 7), (0, 3)), contiguous=False)
-    v1 = View(shape=(4, 2, 6, 2, 1), strides=(42, 21, 3, 1, 0), offset=4, mask=None, contiguous=False)
-    target = View(shape=(4, 2, 6, 2, 1), strides=(2, 1, 12, -4, 0), offset=14, mask=((1, 3), (0, 2), (0, 6), (0, 2), (0, 1)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer4(self):
-    v0 = View(shape=(7, 21, 3), strides=(54, 3, 1), offset=-9, mask=((0, 6), (3, 21), (0, 3)), contiguous=False)
-    v1 = View(shape=(5, 3, 3, 7), strides=(63, 1, 3, 9), offset=63, mask=None, contiguous=False)
-    target = View(shape=(5, 3, 3, 7), strides=(54, 1, 3, 9), offset=45, mask=((0, 5), (0, 3), (0, 3), (1, 7)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer5(self):
-    v0 = View(shape=(5, 1, 24), strides=(20, 0, 1), offset=-2, mask=((0, 5), (0, 1), (2, 22)), contiguous=False)
-    v1 = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 24, 0, 0), offset=0, mask=None, contiguous=False)
-    target = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 20, 0, 0), offset=-2, mask=((1, 11), (0, 2), (0, 5), (0, 2), (0, 1)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  def test_merge_views_variable(self):
-    from tinygrad import Variable
-    N = 100
-    start_pos = Variable("start_pos", 1, N-1)
-    v0 = View(shape=(N, 32, 2), strides=(32, 1, 0), offset=0, mask=((0, N), (0, 32), (0, 1)), contiguous=False)
-    v1 = View(shape=(1, 8, 1, 32), strides=(0, 0, 0, 2), offset=start_pos*64, mask=None, contiguous=False)
-    target = View(shape=(1, 8, 1, 32), strides=(0,0,0,1), offset=start_pos*32, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  def test_view_padded_area1(self):
-    # test_multinomial
-    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
-    v1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1,), strides=(0,), offset=0, mask=((0, 0),), contiguous=False))
-
-  def test_view_padded_area2(self):
-    # test_pad_reflect_mode
-    v0 = View(shape=(1, 1, 10, 7), strides=(0, 0, 5, 1), offset=-15, mask=((0, 1), (0, 1), (3, 8), (0, 5)), contiguous=False)
-    v1 = View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True))
-
-  def test_view_padded_area3(self):
-    # test_roll
-    v0 = View(shape=(2, 4), strides=(0, 1), offset=4, mask=((0, 1), (0, 4)), contiguous=False)
-    v1 = View(shape=(1, 4), strides=(0, 1), offset=4, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 4), strides=(0, 0), offset=0, mask=((0, 0), (0, 0)), contiguous=False))
-
-  def test_view_padded_area4(self):
-    # test_std_mean
-    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((0, 1),), contiguous=False)
-    v1 = View(shape=(1, 1, 1), strides=(0, 0, 0), offset=1, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 1, 1), strides=(0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0)), contiguous=False))
-
-  def test_empty_shape_view1(self):
-    # test_stack_slice
-    v0 = View(shape=(3, 5), strides=(0, 1), offset=0, mask=((0, 1), (0, 5)), contiguous=False)
-    v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(), strides=(), offset=0, mask=None, contiguous=True))
-
-  def test_empty_shape_view2(self):
-    # test_std_mean
-    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
-    v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    # TODO: why is this different?
-    self.assertIsNone(v)
-
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/__init__.py

@@ -15,6 +15,7 @@ from tinygrad.codegen.late.expander import migrate_indexing, expander, pm_pre_ex
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
+from tinygrad.codegen.late.control_flow import pm_add_endrange_endif, pm_control_flow, CFGContext, linearize
 from tinygrad.codegen.opt.postrange import pm_postrange_opt
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_reduce_simplify, pm_flatten_range, pm_split_ranges
 from tinygrad.schedule.rangeify import pm_add_buffers, rangeify_codegen
@@ -101,11 +102,15 @@ def _get_rewrites_for_renderer(opts:Renderer, optimize:bool, linearizer:bool, _Q
   pm_final_rewrite = pm_decomp+pm_render+extra_matcher
   ret.append(RewriteStep(pm_final_rewrite, lambda _: opts.device, name="final rewrite"))
 
-  # return the list (with optional linearizer)
-  return ret + (rewrites_for_linearizer if linearizer else [])
+  # add control flow to the graph
+  ret.append(RewriteStep(pm_add_endrange_endif, name="add endrange/endif"))
+  ret.append(RewriteStep(pm_control_flow, CFGContext, name="add control flow starts", bottom_up=True))
 
-def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True, linearizer:bool=False) -> UOp:
-  return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize, linearizer))
+  # return the list
+  return ret
+
+def full_rewrite_to_sink(sink:UOp, opts:Renderer|None=None, optimize:bool=True) -> UOp:
+  return apply_rewrites(sink, get_rewrites_for_renderer(opts if opts is not None else Renderer(), optimize))
 
 def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
   """
@@ -119,6 +124,6 @@ def full_rewrite(sink:UOp, opts:Renderer|None=None) -> list[UOp]:
     Linear program in UOps.
   """
 
-  lst = list(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None, linearizer=True).arg.lst)
+  lst = linearize(full_rewrite_to_sink(sink, opts, optimize=sink.tag is None))
   if __debug__: type_verify(lst)
   return lst

tinygrad/codegen/gpudims.py

@@ -1,8 +1,7 @@
 import math
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
-from tinygrad.helpers import all_int, dedup
+from tinygrad.helpers import all_int, dedup, get_contraction
 from tinygrad.dtype import dtypes
-from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer
 
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):

tinygrad/codegen/late/control_flow.py

@@ -0,0 +1,103 @@
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
+from tinygrad.helpers import dedup
+from collections import defaultdict
+from itertools import groupby
+from functools import reduce
+import heapq
+
+def linearize(u:UOp) -> list[UOp]:
+  lst = list(u.toposort())
+  in_this_block = set(lst)
+  local_children: defaultdict[UOp, list[UOp]] = defaultdict(list)
+  in_degree:dict[UOp, int] = {}
+  priorities:dict[UOp, int] = {}
+
+  # get local children and assign priorities
+  # NOTE: this requires the lst be locally toposorted
+  for u in reversed(lst):
+    in_degree[u] = 0
+    for s in u.src:
+      if s in in_this_block:
+        local_children[s].append(u)
+        in_degree[u] += 1
+    # put loads in the beginning of the block and prevent priority inversion. hack for BARRIER grouping too
+    priority = [0] + [priorities[x] for x in local_children[u]]
+    if u.op is Ops.LOAD: priority.append(-1000)
+    if u.op is Ops.BARRIER: priority.append(-1500)
+    # ranges are scheduled as late as possible so anything that can be outside is
+    #if u.op is Ops.RANGE: priority = [2000]
+    # move defines and consts to the top
+    if u.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}: priority.append(-2000)
+    priorities[u] = min(priority)
+
+  # number the uops in "ideal" order
+  nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: (priorities[x],)+x.tuplize))}
+
+  # then force then to be toposorted in as close to the ideal order as possible
+  heapq.heapify(heap:=[(nkey[u],u) for u in lst if in_degree[u] == 0])
+  newlst = []
+  while heap:
+    newlst.append(u:=heapq.heappop(heap)[1])
+    for v in local_children[u]:
+      in_degree[v] -= 1
+      if in_degree[v] == 0: heapq.heappush(heap, (nkey[v],v))
+
+  assert len(newlst) == len(lst), f"len mismatch {len(newlst)} != {len(lst)}"
+  return newlst
+
+def add_endrange(x:UOp):
+  if not ((x.op is Ops.LOAD and x.src[-1].op is Ops.STORE) or all(s.op is Ops.STORE and any(n.op is Ops.RANGE for n in s.src) for s in x.src)):
+    return None
+  src: list[UOp] = []
+  for k,g in groupby(x.src, key=lambda k: tuple(dedup(s for s in k.src if s.op is Ops.RANGE))):
+    if not k: src.extend(g)
+    else: src.extend(reduce(lambda acc,rng: (UOp(Ops.ENDRANGE, src=(rng,) + acc),), reversed(k), tuple(g))) # type: ignore
+  return x.replace(src=tuple(src))
+
+def add_endif(x:UOp):
+  groups = {k: tuple(g) for k,g in groupby(x.src, key=lambda k: k.src[2] if len(k.src) >= 3 and k.src[2].op is Ops.IF else k)}
+  if not any(k.op is Ops.IF for k in groups): return None
+  return x.replace(src=tuple(UOp(Ops.ENDIF, src=(k,) + g) if k.op is Ops.IF else k for k,g in groups.items()))
+
+# some Ops.IF aren't closed by an Ops.STORE, in that case the Ops.SINK closes it
+def close_ifs(x:UOp):
+  consumers = x.get_consumer_map()
+  if (y:=next((s for s in consumers if s.op is Ops.IF and all(n.op is not Ops.ENDIF for n in consumers[s])), None)) is not None:
+    return x.replace(src=(UOp(Ops.ENDIF, src=(y,) + x.src),))
+  return None
+
+pm_add_endrange_endif = PatternMatcher([
+  (UPat((Ops.SINK, Ops.NOOP, Ops.LOAD), name="x"), add_endrange),
+  (UPat((Ops.SINK, Ops.ENDRANGE, Ops.BARRIER), name="x"), add_endif),
+  (UPat(Ops.SINK, name="x"), close_ifs),
+])
+
+class CFGContext:
+  def __init__(self, sink:UOp):
+    # there are 3 relationships between ranges:
+    # nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
+    # dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
+    # independent, endrange y is not a dependency of endrange x
+    deps: dict[UOp, set[UOp]] = {}
+    nesting: dict[UOp, UOp] = {}
+    for u in sink.toposort():
+      deps[u] = set().union(*(deps[s] for s in u.src))
+      if u.op in (Ops.ENDRANGE, Ops.ENDIF):
+        for n in [x for x in deps[u] if x.op in (Ops.ENDRANGE, Ops.ENDIF) and u.src[0] in deps[x] and x not in nesting]: nesting[n] = u
+      if u.op is Ops.SINK:
+        for n in [x for x in deps[u] if x.op in (Ops.ENDRANGE, Ops.ENDIF) and x not in nesting]: nesting[n] = u
+      if u.op in (Ops.RANGE, Ops.ENDRANGE, Ops.IF, Ops.ENDIF): deps[u] |= {u}
+
+    self.edges: dict[UOp, UOp] = {}
+    siblings: dict[UOp, list[UOp]] = {}
+    for k,vv in nesting.items(): siblings.setdefault(vv, []).append(k)
+    for k,v in siblings.items():
+      # range/if that have dependencies on other siblings need to run after them
+      order = sorted(v, key=lambda x: len([y for y in v if y in deps[x]]))
+      zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[0]] + order, order)
+      for x,y in zipped: self.edges[y.src[0]] = x
+
+pm_control_flow = PatternMatcher([
+  (UPat((Ops.RANGE, Ops.IF), src=(UPat(),), name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
+  (UPat(Ops.IF, src=(UPat(), UPat(Ops.BARRIER)), name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
+])

tinygrad/codegen/late/devectorizer.py

@@ -11,7 +11,7 @@ from tinygrad.renderer import Renderer
 # ***** image load valid simplification *****
 
 def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
-  if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.index(UOp.invalid())
+  idx = uop_given_valid(valid, start_idx)
   if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid))
 
   # wait for it to be image indexed before running simplification

tinygrad/helpers.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass
-import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal
+import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator
 
@@ -82,6 +82,13 @@ def word_wrap(x, wrap=80):
   while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
   return x[:i] + "\n" + word_wrap(x[i:], wrap)
 
+# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
+def get_contraction(old_shape:tuple[T, ...], new_shape:tuple[T, ...]) -> list[list[int]]|None: # T is sint
+  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
+  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
+  except ValueError: return None
+  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
+
 def suppress_finalizing(func):
   def wrapper(*args, **kwargs):
     try: return func(*args, **kwargs)

tinygrad/shape/view.py

@@ -4,14 +4,7 @@ from dataclasses import dataclass
 from typing import cast, Sequence
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import resolve, UOp, Variable, sint, smax, smin, sint_to_uop, Ops, ssimplify
-from tinygrad.helpers import prod, all_int, flatten, ceildiv
-
-# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
-def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
-  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
-  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
-  except ValueError: return None
-  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
+from tinygrad.helpers import prod, all_int, flatten
 
 @functools.cache
 def canonicalize_strides(shape:tuple[sint, ...], strides:tuple[sint, ...]) -> tuple[sint, ...]:
@@ -171,7 +164,6 @@ class View:
     if not all_int(vm1.shape):
       # if all strides are 0 and vm2 is unmasked, return vm1
       if all(x == 0 for x in vm2.strides+vm1.strides) and vm2.mask is None: return vm1
-      # TODO: handle more cases
       return None
 
     # Project vm1's offset and strides on to vm2.
@@ -184,47 +176,7 @@ class View:
         if not resolve((s1 := s1 - o)!=0): continue  # if s1 can possibly be 0
         terms[d2].append((d1, s1))
         strides[d1] += ssimplify(s1 * vm2.strides[d2])
-
-    # Merge dimensions in vm2 if required.
-    # NB: Merging too many dimensions can make it difficult to project vm2's mask, hence only combining when required.
-    idxs: list[UOp] = [UOp.variable(f"idx{i}", 0, s-1, dtypes.index) for i,s in enumerate(vm1.shape)]
-    merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
-    extents: list[tuple[sint, UOp]] = []
-    for term, s, o in zip(reversed(terms), reversed(vm2.shape), reversed(origin)):
-      merged_term += (sum([idxs[d1] * s1 for d1, s1 in term]) + o) * merged_size
-      merged_size *= s
-      if resolve(merged_term < merged_size, False) and resolve(0 <= merged_term, False):
-        extents.append((merged_size, merged_term))
-        merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
-    if resolve(merged_term != 0): return None
-    if (vm2_shape := tuple(s for s,_ in reversed(extents))) != vm2.shape:
-      if (reshaped_vm2 := vm2.reshape(vm2_shape)) is None: return None
-      # NOTE: this != to prevent infinite loop
-      if reshaped_vm2.shape != vm2.shape: return reshaped_vm2 + vm1
-
-    if vm2.mask:
-      # Try to project vm2's mask on to vm1.
-      newb, newe, bad = [0] * len(vm1.shape), list(vm1.shape), False
-      for (b, e), o, term, (_, t) in zip(vm2.mask, origin, terms, reversed(extents)):
-        if resolve(b <= (t := t.simplify()).vmin and t.vmax < e, False): continue
-        if len(term) != 1:
-          if not term and newe:
-            # t should be a constant if no terms contribute to this dimension, but it might not be simplified
-            if t.vmin != t.vmax: return None
-            newe[0] = 0
-          else: bad = True
-          continue
-        d1, s1 = term[0]
-        newb[d1] = smax(newb[d1], ceildiv(b - o if s1 > 0 else e - o - 1, s1))
-        newe[d1] = smin(newe[d1], (b - o if s1 < 0 else e - o - 1) // s1 + 1)
-
-      # If any of vm1 was masked off, try again with that mask in place.
-      if any((b, e) != (0, s) for b, e, s in zip(newb, newe, vm1.shape)):
-        return vm2 + View.create(vm1.shape, vm1.strides, vm1.offset, tuple(zip(newb, newe)))
-      # Otherwise if vm2's mask was violated, then cannot merge.
-      if bad: return None
-
-    return View.create(vm1.shape, tuple(strides), ssimplify(sum(o * s for o, s in zip(origin, vm2.strides)) + vm2.offset))
+    return None
 
   def __unsafe_resize(self, arg: tuple[tuple[sint, sint], ...], mask=None) -> View:
     offset = sum([s * x[0] for s, x in zip(self.strides,arg)])
@@ -292,7 +244,6 @@ class View:
 
     r_strides, r_new_shape = [], reversed(new_shape)
     for merged_size, new_stride, real_size in reversed(merge_dims(self.shape, self.strides, self.mask)):
-      # TODO: write with get_contraction
       acc = 1
       # TODO: third resolve shouldn't be needed
       while resolve(acc <= merged_size) and resolve(acc != merged_size) and resolve((new_dim := next(r_new_shape, 0)) > 0):

tinygrad/uop/ops.py

@@ -240,6 +240,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
         if s in ret: del ret[s]
     else:
       for s in self.src: ret.update(s.ranges)
+    if self.op is Ops.ENDRANGE: del ret[self.src[0]]
     return ret
 
   @property

tinygrad/uop/spec.py

@@ -153,7 +153,8 @@ spec = PatternMatcher([
   (UPat(Ops.DEFINE_REG, src=()), lambda: True),
   (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)),
 
-  (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
+  (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng", allow_any_len=True),
+   lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \
      all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)),
   (UPat(Ops.SPECIAL, src=(UPat.var("x"),), name="s"), lambda s,x: s.dtype == x.dtype == dtypes.int32 and isinstance(s.arg, str)),
 
@@ -190,7 +191,7 @@ spec = PatternMatcher([
   (UPat((Ops.IDIV, Ops.MOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
   (UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype.base == y.dtype.base for y in x.src)),
 
-  (UPat(Ops.ENDRANGE, dtype=dtypes.void, src=(UPat(Ops.RANGE),)), lambda: True),
+  (UPat(Ops.ENDRANGE, dtype=dtypes.void, src=(UPat(Ops.RANGE),), allow_any_len=True), lambda: True),
 
   # WMMA has a <a, b, acc>
   (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 8),
@@ -198,9 +199,8 @@ spec = PatternMatcher([
   (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
 
   # if has a <gate, barrier?>
-  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(),)), lambda: True),
-  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(), UPat(Ops.BARRIER))), lambda: True),
-  (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True),
+  (UPat(Ops.IF, dtype=dtypes.void, src=(UPat(),), allow_any_len=True), lambda: True),
+  (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),), allow_any_len=True), lambda: True),
 
   (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) >= 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}),
   (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()),

tinygrad/uop/symbolic.py

@@ -397,8 +397,8 @@ def parse_valid(valid:UOp) -> tuple[UOp, bool, int]:
   if valid.op is Ops.CMPLT and dtypes.is_int(valid.src[0].dtype): return valid.src[0], True, int((valid.src[1]).vmax)-1
   raise ValueError(f"not able to parse {valid=}")
 
-def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
-  # return None if valid is always False, otherwise the simplified uop (might be the same as input)
+def uop_given_valid(valid:UOp, uop:UOp) -> UOp:
+  # return simplified uop (might be the same as input)
 
   # first, parse valid into {expr: (lower_bound, upper_bound)}
   bounds:defaultdict[UOp, list[ConstType|None]] = defaultdict(lambda: [None, None])
@@ -415,18 +415,13 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
     v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
     expr = expr.substitute(load_subs)  # make sure expr appears in same form in the uop
     # some expr has lower bound > upper bound -> valid is an empty set and we return None
-    if v0 > v1: return None
-    # whole node became a const
-    if v0 == v1:
-      uop = uop.substitute({expr:expr.const_like(v0)}).simplify()
-      continue
     # every candidate is a set of constrained UOp based on valid, and if every item in a set simplifies the uop into a same output, we rewrite uop
     candidates = []
     if expr.op is Ops.ADD and v0 == 1 and all(u.op in GroupOp.Irreducible for u in expr.split_uop(Ops.ADD)):
       # if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output
       candidates.append([(Xi, UOp.variable("fake", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)])
     # try checking the whole clause
-    if expr in uop.toposort(): candidates.append([(expr, UOp.variable("fake", v0, v1, expr.dtype))])
+    candidates.append([(expr, UOp.variable("fake", v0, v1, expr.dtype))])
 
     for candidate in candidates:
       # if every branch in candidate gives the same simplified uop, we can rewrite the uop
@@ -451,7 +446,7 @@ def simplify_valid(valid:UOp) -> UOp|None:
   something_changed = False
   valids = list(valid.split_uop(Ops.AND))
   for stmt in sorted(valids, key=lambda v: _valid_priority(v, valids)):
-    ret.append(newstmt if ret and (newstmt:=uop_given_valid(functools.reduce(operator.and_, ret), stmt)) is not None else stmt)
+    ret.append(uop_given_valid(functools.reduce(operator.and_, ret), stmt) if ret else stmt)
     if ret[-1] is not stmt: something_changed = True
   return functools.reduce(operator.and_, ret) if something_changed else None
 

tinygrad/viz/index.html

@@ -315,15 +315,6 @@
     font-size: 0.95em;
     letter-spacing: 0.03em;
   }
-  .legend {
-    display: flex;
-    align-items: center;
-  }
-  .legend > div {
-    width: 0.95em;
-    height: 0.95em;
-    margin-right: 4px;
-  }
   </style>
 </head>
 <body>

tinygrad/viz/js/index.js

@@ -164,6 +164,12 @@ const drawLine = (ctx, x, y, opts) => {
   ctx.stroke();
 }
 
+function tabulate(rows) {
+  const root = d3.create("div").style("display", "grid").style("grid-template-columns", `${Math.max(...rows.map(x => x[0].length), 0)}ch 1fr`).style("gap", "0.2em");
+  for (const [k,v] of rows) { root.append("div").text(k); root.append("div").node().append(v); }
+  return root;
+}
+
 var data, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
 async function renderProfiler() {
   displayGraph("profiler");
@@ -272,7 +278,10 @@ async function renderProfiler() {
       for (const [num, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) {
         const x = steps.map(s => timestamps[s]);
         const dur = x.at(-1)-x[0];
-        const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}\nnum:${num}\nalive for ${formatTime(dur)}`};
+        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};
         shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
       }
       // generic polygon merger
@@ -434,6 +443,7 @@ async function renderProfiler() {
     e.preventDefault();
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
     if (foundRect?.step != null) return setCtxWithHistory(foundRect.ctx, foundRect.step);
+    return document.querySelector(".metadata").replaceChildren(foundRect?.html ?? "");
   });
 
   canvas.addEventListener("mousemove", e => {
@@ -644,17 +654,10 @@ async function main() {
           }
         }
       }
-      const summary = metadata.appendChild(document.createElement("table"));
-      for (const s of ret.summary) {
-        const tr = summary.appendChild(document.createElement("tr"));
-        tr.className = "main-row";
-        const td = tr.appendChild(document.createElement("td"));
-        const div = td.appendChild(document.createElement("div"));
-        div.className = "legend";
-        div.appendChild(document.createElement("div")).style.background = cycleColors(colorScheme.CATEGORICAL, s.idx);
-        div.appendChild(document.createElement("p")).textContent = s.label;
-        appendTd(tr, s.value);
-      }
+      metadata.appendChild(tabulate(ret.summary.map(s => {
+        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "24px").style("height", "100%");
+        return [s.label.trim(), div.node()];
+      })).node());
     } else root.appendChild(codeBlock(ret.src, "x86asm"));
     return document.querySelector(".disasm").replaceChildren(root);
   }

tinygrad/viz/serve.py

@@ -206,7 +206,7 @@ def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
   # disassembly output can include headers / metadata, skip if llvm-mca can't parse those lines
   data = json.loads(subprocess.check_output(["llvm-mca","-skip-unsupported-instructions=parse-failure","--json","-"]+target_args, input=asm.encode()))
   cr = data["CodeRegions"][0]
-  resource_labels = data["TargetInfo"]["Resources"]
+  resource_labels = [repr(x)[1:-1] for x in data["TargetInfo"]["Resources"]]
   rows:list = [[instr] for instr in cr["Instructions"]]
   # add scheduler estimates
   for info in cr["InstructionInfoView"]["InstructionList"]: rows[info["Instruction"]].append(info["Latency"])