diff --git a/test/null/test_gpudims.py b/test/null/test_gpudims.py index 0a6fb27037..a9a3d10ee0 100644 --- a/test/null/test_gpudims.py +++ b/test/null/test_gpudims.py @@ -94,6 +94,17 @@ class TestGroupedDims(unittest.TestCase): assert idxs[2].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[2].op}" assert idxs[3].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[3].op}" + def test_grouped_dims_high_rank(self): + # 4D collapsed onto 2 axes + self._check_grouped_dims("gidx", (4,4,4,4), (16,16), False, [16,16]) + # 4D untouched + self._check_grouped_dims("gidx", (2,3,4,5), None, False, [2,3,4,5]) + idxs = get_grouped_dims("gidx", (2,3,4,5), None, False) + assert all(u.op is Ops.SPECIAL for u in idxs), f"expected all-SPECIAL when untouched, got {[u.op for u in idxs]}" + # 5D and 6D collapsed onto 3 axes + self._check_grouped_dims("gidx", (2,2,2,2,2), (4,4,4), False, [4,4,2]) + self._check_grouped_dims("gidx", (2,2,2,2,2,2), (8,8,8), False, [8,4,2]) + def test_global_prod_max(self): g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL) sink = UOp.param(0, dtypes.float.ptr()).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo()) diff --git a/test/null/test_helpers.py b/test/null/test_helpers.py index 44c49e8891..d377cc34c4 100644 --- a/test/null/test_helpers.py +++ b/test/null/test_helpers.py @@ -1,6 +1,6 @@ import ctypes, gzip, unittest, timeit, pickle from tinygrad import Variable -from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count, all_same +from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, count, all_same 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.helpers import ceildiv, ansistrip, get_shape from tinygrad.tensor import Tensor @@ -273,75 +273,6 @@ class TestMemoryview(unittest.TestCase): mva_us = timeit.timeit(lambda: mv_address(x), number=iters) * 1e6 / iters print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs") -class TestGetContraction(unittest.TestCase): - def test_contraction(self): - r = get_contraction((1,2,3,4), (2,3,4)) - self.assertEqual(r, [[0, 1], [2], [3]]) - - r = get_contraction((2,1,3,4), (2,3,4)) - self.assertEqual(r, [[0], [1, 2], [3]]) - - r = get_contraction((1,2,3,1,4), (1,2,3,4)) - self.assertEqual(r, [[], [0, 1], [2], [3, 4]]) - - r = get_contraction((1,2,3,1,4,1,1), (2,3,4)) - self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]]) - - r = get_contraction((1,2,3,4), (1,2,3*4)) - self.assertEqual(r, [[], [0, 1], [2, 3]]) - - r = get_contraction((1,2,3,4), (2,1,3,4)) - self.assertEqual(r, [[0, 1], [], [2], [3]]) - - r = get_contraction((1,2,3,4), (1,1,2*3*4,1)) - self.assertEqual(r, [[], [], [0,1,2,3], []]) - - r = get_contraction((2,1,3,4), (1,2,3,4)) - self.assertEqual(r, [[], [0], [1, 2], [3]]) - - r = get_contraction((1,2,3,4), (2*3*4,1,1,1)) - self.assertEqual(r, [[0, 1, 2, 3], [], [], []]) - - r = get_contraction((4,4,4,4), (16,1,16)) - self.assertEqual(r, [[0, 1], [], [2, 3]]) - - r = get_contraction((1,2,3,4,1,1,1), (2,3,4)) - self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]]) - - r = get_contraction((1,2,3,4), (1,2,3,4,1)) - self.assertEqual(r, [[], [0, 1], [2], [3], []]) - - r = get_contraction((14,1,384,14,1,1,1,1), (1,14,384,14)) - self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7]]) - - r = get_contraction((14,1,384,1,14,1,1,1,1), (1,14,384,14)) - self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7,8]]) - - r = get_contraction((512, 512), (1, 1, 512, 1, 1, 1, 1, 512)) - self.assertEqual(r, [[], [], [0], [], [], [], [], [1]]) - - r = get_contraction((1,2,3,4), (1,2,6,2)) - self.assertEqual(r, None) - - def test_contraction_ones(self): - r = get_contraction((1,), (1,1,1)) - self.assertEqual(r, [[], [], [0]]) - - r = get_contraction((1,1), (1,1,1)) - self.assertEqual(r, [[], [], [0, 1]]) - - r = get_contraction((1,1,1,1), (1,)) - self.assertEqual(r, [[0,1,2,3]]) - - r = get_contraction((1,1,1,1), (1,1)) - self.assertEqual(r, [[], [0,1,2,3]]) - - r = get_contraction((1,1,1,1), (1,1,1)) - self.assertEqual(r, [[], [], [0,1,2,3]]) - - r = get_contraction((1,1,1,1), (1,1,1,1)) - self.assertEqual(r, [[], [], [], [0,1,2,3]]) - class TestGetShape(unittest.TestCase): def test_get_shape(self): assert get_shape(2) == () diff --git a/tinygrad/codegen/gpudims.py b/tinygrad/codegen/gpudims.py index 7ffd44c2cc..4688470e8f 100644 --- a/tinygrad/codegen/gpudims.py +++ b/tinygrad/codegen/gpudims.py @@ -1,6 +1,6 @@ import math from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType -from tinygrad.helpers import dedup, get_contraction +from tinygrad.helpers import dedup from tinygrad.dtype import dtypes, AddrSpace, Invalid from tinygrad.renderer import Renderer @@ -36,24 +36,8 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No # try to split up dims: (a,) -> (b, c) if limited == dims: limited = _split_dims(dims, max_sizes) raw_idxs = [UOp.special(s, f"{prefix}{i}") for i,s in enumerate(limited)] - if len(limited) < len(dims): - ret = [] - if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(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) - return ret - elif (a:=len(limited)) > (b:=len(dims)): - if a == 2 and b == 1: return [raw_idxs[0] * limited[1] + raw_idxs[1]] - if a == 3 and b == 1: return [(raw_idxs[0] * limited[1] + raw_idxs[1]) * limited[2] + raw_idxs[2]] - if limited != dims: - # Convert to 1D - flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(limited) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2] - # Get back original indices from 1D - return [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]] - return raw_idxs + flat = sum(idx * math.prod(limited[i+1:]) for i,idx in enumerate(raw_idxs)) + return [ssimplify(flat // math.prod(dims[i+1:])) if i == 0 else ssimplify((flat // math.prod(dims[i+1:])) % dims[i]) for i in range(len(dims))] def add_gpudims(ctx:Renderer, s:UOp): if s.arg is None: return None diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index fe30002c50..a6e6def4d8 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -122,13 +122,6 @@ def strides_for_shape(shape:tuple[T, ...]) -> tuple[T, ...]: strides = tuple(itertools.accumulate(reversed(shape[1:]), operator.mul, initial=1))[::-1] return canonicalize_strides(shape, strides) -# 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 = [0 if isinstance(acc, int) and acc == 1 else acc_old.index(acc)+1 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)