okay to skip that test

* try removing factors of two

* dont return if None

* add test

docs/tinybox.md

@@ -6,7 +6,7 @@ If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygra
 
 ## Welcome
 
-Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, and the green box includes six 4090 GPUs. Whether you bought a red one or a green one, we want you to love it.
+Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.
 
 We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.
 

setup.py

@@ -29,6 +29,7 @@ setup(name='tinygrad',
         'tinygrad.apps',
         'tinygrad.codegen',
         'tinygrad.codegen.opt',
+        'tinygrad.codegen.late',
         'tinygrad.engine',
         'tinygrad.frontend',
         'tinygrad.nn',

test/device/test_hcq.py

@@ -155,6 +155,7 @@ class TestHCQ(unittest.TestCase):
     val = TestHCQ.b.uop.buffer.as_buffer().cast("f")[1]
     assert val == 0.0, f"got val {val}, should not be updated"
 
+  @unittest.skip("globals/locals are merged now")
   @unittest.skipIf(Device.DEFAULT in {"CPU", "LLVM"}, "No globals/locals on LLVM/CPU")
   def test_exec_update_fuzz(self):
     virt_val = Variable("sig_val", 0, 0xffffffff, dtypes.uint32)

test/test_linearizer.py

@@ -3,7 +3,6 @@ import unittest
 from dataclasses import replace
 
 from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
-from tinygrad.codegen.gpudims import get_grouped_dims
 from tinygrad.uop.ops import UOp, Ops, GroupOp, KernelInfo
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -467,77 +466,7 @@ class TestLinearizer(unittest.TestCase):
         end_range = [i for i, x in enumerate(uops) if x.op is Ops.ENDRANGE][0]
         assert end_range < uops.index(u)
 
-  def test_grouped_dims(self):
-    def _assert_grouped_dims(prefix, dims, max_sizes, reverse_dims, expected_sizes, assert_same_length = True):
-      idxs = get_grouped_dims(prefix, dims, max_sizes, reverse_dims)
-      loop_idxs = dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs]))
-      loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg[0])
-      sizes = [x.arg[1] for x in loop_idxs]
-      assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}"
-      if assert_same_length:
-        assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}"
-      assert sizes == expected_sizes, f"expected sizes={expected_sizes}, got {sizes=}"
-      # TODO: add these back after uop symbolic
-      # for i in range(len(dims)):
-      #   assert idxs[i].max+1 == dims[i], f"idxs[{i}] should have max {dims[i]-1}"
-      # for i in range(len(loop_idxs)):
-      #   assert loop_idxs[i].expr.startswith(prefix), f"loop_idxs[{i}] must start with {prefix}"
-      #   assert loop_idxs[i].max+1 == sizes[i], f"loop_idxs[{i}] should have max {sizes[i]-1}"
-
-    # no-op
-    _assert_grouped_dims("gidx", (2,), (16,16,16), False, [2])
-    _assert_grouped_dims("gidx", (2,3), (16,16,16), False, [2,3])
-
-    # check reverse dims
-    _assert_grouped_dims("gidx", (2,3), (16,16,16), True, [3,2])
-    _assert_grouped_dims("gidx", (2,3,4), (16,16,16), False, [2,3,4])
-
-    # test splitting globals:    len(dims) == len(max)
-    _assert_grouped_dims("gidx", (64,3,4), (16,16,16), False, [16,12,4])
-    _assert_grouped_dims("gidx", (64,3,4), (16,4,16), False, [16,3,16])
-    _assert_grouped_dims("gidx", (64,3,4), (16,16,16), True, [16,3,16])
-    _assert_grouped_dims("gidx", (128,3,4), (16,4,256), False, [16,3,32])
-    _assert_grouped_dims("gidx", (4,4,512), (16,4,256), False, [8,4,256])
-
-    # prefer group_dim strategy when possible
-    _assert_grouped_dims("gidx", (512,4,2), (8192,2,2), False, [2048,2])
-
-    # test splitting globals:    len(dims) < len(max)
-    #                            len(dim)        ->          len(limited)
-    #                              1             ->             2
-    _assert_grouped_dims("gidx", (128,), (16,16,256), False, [16,8], False)
-    #                              1             ->             3
-    _assert_grouped_dims("gidx", (65536,), (16,16,256), False, [16,16,256], False)
-    #                              2             ->             3
-    _assert_grouped_dims("gidx", (128,128), (16,16,256), False, [16,16,64], False)
-    # test when the only divisor is the square root of dim
-    _assert_grouped_dims("gidx", (121,), (12,12,12), False, [11,11], False)
-
-    # collapse on onto the left most axis
-    _assert_grouped_dims("gidx", (2,3,4,5), (16,16,16), False, [6,4,5])
-    _assert_grouped_dims("gidx", (2,3,4,5), (32,16,16), True, [20,3,2])
-    # _assert_grouped_dims("gidx", (Variable("start_pos",1,2),3,4,5), (32,16,16), True, [20,3,Variable("start_pos",1,2)])
-
-    # collapse on left-most available axis (the left most is too small)
-    _assert_grouped_dims("gidx", (2,3,4,5), (4,16,16), False, [2,12,5])
-    _assert_grouped_dims("gidx", (2,3,4,5), (16,16,16), True, [5,12,2])
-
-    # _assert_grouped_dims("gidx", (Variable("start_pos",1,2),3,4,5), (16,16,16), False, [Variable("start_pos",1,2)*3,4,5])
-
-    # dim too large and not factorable
-    with self.assertRaises(RuntimeError):
-      get_grouped_dims("gidx", (23,), (16,16,16), False,)
-    with self.assertRaises(RuntimeError):
-      get_grouped_dims("gidx", (128,3,4), (16,2,2), False,)
-
-    # too large for sizes
-    with self.assertRaises(RuntimeError):
-      get_grouped_dims("gidx", (2,3,4,5,6), (16,16,16))
-
-    # # variable too large
-    # with self.assertRaises(AssertionError):
-    #   get_grouped_dims("gidx", (Variable("start_pos",0,16),3,4), (16,16,16), False,)
-
+  @unittest.skip("only one global now")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   def test_default_global_reversed(self):
     # shrink so that the dims do not collapse

test/test_pickle.py

@@ -20,7 +20,7 @@ class TestPickle(unittest.TestCase):
     self.assertEqual(pm2.rewrite(sink).key, tt.key)
 
   def test_pickle_main_pattern_matcher(self):
-    from tinygrad.codegen.devectorizer import sym
+    from tinygrad.codegen.late.devectorizer import sym
     ssym = pickle.dumps(sym)
     dsym = pickle.loads(ssym)
     self.assertEqual(dsym.patterns[0][0].location, sym.patterns[0][0].location)

test/test_uop_graph.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
 from tinygrad.uop.symbolic import sym
 from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
-from tinygrad.codegen.expander import expander
+from tinygrad.codegen.late.expander import expander
 
 simple_pm = PatternMatcher([
   (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),

test/test_uops.py

@@ -402,6 +402,14 @@ class TestAssembly(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.IDIV, ops)
 
+  def test_fast_idiv_remove_powers_of_two(self):
+    ridx = UOp.range(dtypes.int, 2**20, 0)
+    uops = to_uops_list([ridx//(7*64)], opts=Device[Device.DEFAULT].renderer)
+    ops = [x.op for x in uops]
+    # this requires shifting out the powers of two before doing fast_idiv
+    # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
+    self.assertNotIn(Ops.CAST, ops)
+
   def test_mulacc_unrolled(self):
     # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
     # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)

test/unit/test_block_reorder.py

@@ -1,7 +1,7 @@
 import unittest, random
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import print_uops, UOp, Ops
-from tinygrad.codegen.linearize import block_reorder
+from tinygrad.codegen.late.linearize import block_reorder
 from tinygrad.renderer.cstyle import OpenCLRenderer
 
 def is_toposorted(lst:list[UOp]):

test/unit/test_shapetracker.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import prod
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad import Variable
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.codegen.devectorizer import sym
+from tinygrad.codegen.late.devectorizer import sym
 from itertools import product
 
 def shapetracker_getitem(st:ShapeTracker, val:int):

test/unit/test_uop_symbolic.py

@@ -4,7 +4,7 @@ import z3
 
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.codegen import full_rewrite
-from tinygrad.codegen.devectorizer import sym
+from tinygrad.codegen.late.devectorizer import sym
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad import Variable

tinygrad/codegen/__init__.py

@@ -12,10 +12,10 @@ from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.uop.decompositions import get_late_rewrite_patterns
-from tinygrad.codegen.expander import migrate_indexing, expander
-from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
+from tinygrad.codegen.late.expander import migrate_indexing, expander
+from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
+from tinygrad.codegen.late.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 from tinygrad.codegen.opt import pm_get_optimization, pm_do_optimize
 from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
 

tinygrad/codegen/gpudims.py

@@ -1,53 +1,15 @@
-import math
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
-from tinygrad.helpers import all_int, partition, flatten, prod, dedup
+from tinygrad.helpers import partition, flatten, prod, dedup
 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, ...]):
-  # TODO: symbolic shape
-  if not all_int(dims): return dims
-  while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
-    for i,m in enumerate(max_sizes):
-      if i < (len(dims)-1) and dims[i] * dims[i+1] <= m:
-        dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
-        break
-    else: return None
-  return dims
-
-def _split_dims(dims, max_sizes):
-  if all(d <= m for d,m in zip(dims, max_sizes)): return dims
-  _dims = list(dims) + [1]*(3-len(dims))
-  for i in range(len(_dims)):
-    while _dims[i] > max_sizes[i]:
-      div = next((d for d in range(2, math.ceil(math.sqrt(_dims[i])) + 1) if (_dims[i] % d) == 0), 1)
-      if div == 1: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
-      _dims[i], _dims[(i+1)%len(_dims)] = _dims[i]//div, _dims[(i+1)%len(_dims)]*div
-  return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
-
-def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
+def get_grouped_dims(prefix, dims:tuple[sint, ...], reverse=False) -> list[UOp]:
   if reverse: dims = dims[::-1]
-  # try to group first: (a, b, c, d) -> (ab, c, d)
-  limited = (grouped if (grouped := _group_dims(dims, max_sizes)) else dims) if max_sizes is not None else dims
-  # check if grouping failed
-  if max_sizes is not None and len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
-  # try to split up dims: (a,) -> (b, c)
-  if limited == dims: limited = _split_dims(dims, max_sizes) if max_sizes is not None else dims
-  ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (), (f"{prefix}{i}", s)) for i,s in enumerate(limited)]
-  if len(limited) < len(dims):
-    ret = []
-    if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
-    for idx, contraction_group in zip(raw_idxs, contraction):
-      for c in contraction_group[:-1]:
-        ret.append(idx % dims[c])
-        idx //= dims[c]
-      ret.append(idx)
-  elif len(limited) > len(dims):
-    a, b = len(limited), len(dims)
-    if a == 2 and b == 1: ret = [raw_idxs[0] * limited[1] + raw_idxs[1]]
-    if a == 3 and b == 1: ret = [raw_idxs[0] * (limited[1] * limited[2]) + raw_idxs[1] * limited[2] + raw_idxs[2]]
-    if a == 3 and b == 2: ret = [raw_idxs[0] * limited[1] + raw_idxs[1], raw_idxs[2]]
+  spec = UOp(Ops.SPECIAL, dtypes.int, (), (f"{prefix}0", ssimplify(prod(dims))))
+  ret = []
+  for d in dims:
+    ret.append(spec % d)
+    spec //= d
   return ret[::-1] if reverse else ret
 
 def add_gpudims(ctx:Renderer, s:UOp):
@@ -72,10 +34,10 @@ def add_gpudims(ctx:Renderer, s:UOp):
   ki: KernelInfo = s.arg
   if ki.dont_use_locals:
     assert not local_dims, "can't use locals if there's no local dims"
-    idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
+    idxs = get_grouped_dims("idx", global_shape, reverse=True)
   else:
     # define indexes for GPU-like execution
-    idxs = get_grouped_dims("gidx", global_shape, ctx.global_max, reverse=True) + get_grouped_dims("lidx", local_shape, ctx.local_max)
+    idxs = get_grouped_dims("gidx", global_shape, reverse=True) + get_grouped_dims("lidx", local_shape)
 
   # apply to multiple ranges
   subs = {}

tinygrad/renderer/cstyle.py

@@ -6,7 +6,7 @@ from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
 from tinygrad.renderer import Renderer
-from tinygrad.codegen.devectorizer import no_vectorized_alu
+from tinygrad.codegen.late.devectorizer import no_vectorized_alu
 
 base_rewrite = PatternMatcher([
   (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),

tinygrad/uop/decompositions.py

@@ -280,7 +280,7 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]:
       return m, s
   assert False
 
-def fast_idiv(device: str, x: UOp, d: int) -> UOp|None:
+def fast_idiv(device: str, x: UOp, d: int, dont_cast=False) -> UOp|None:
   # If d is a power of two this is not valid for signed ints!
   is_unsigned = True if x.vmin>=0 or x.dtype in dtypes.uints else False
   assert d>0, "Sign should have been taken out of divisor"
@@ -288,6 +288,10 @@ def fast_idiv(device: str, x: UOp, d: int) -> UOp|None:
   m,s = magicgu(max(vmax, abs(vmin)), d)
   if m*vmin >= dtypes.min(x.dtype) and m*vmax <= dtypes.max(x.dtype):
     return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
+  # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
+  if (largest_factor_of_two_in_d := (d & -d)) > 1:
+    if (ret:=fast_idiv(device, x//largest_factor_of_two_in_d, d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
+  if dont_cast: return None
   # promo_lattice needs to return an unsigned type if the type is unsigned
   if dtypes.is_int(next_dtype := promo_lattice[x.dtype][-1]) and is_dtype_supported(next_dtype, None if device=='' else device):
     if m*vmin >= dtypes.min(next_dtype) and m*vmax <= dtypes.max(next_dtype):