From 36ce883c7d691f4d385a5bb12d4364a25e0c82d0 Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 14 Jul 2025 11:10:30 -0400 Subject: [PATCH] update heuristic to use k.upcastable_dims and k.unrollable_dims [pr] (#11233) idea is to make it behave the same regardless of axis order and with empty 1s in shape. not quite fully remove all first_upcast yet because some conditions used already upcasted size which need a separate benchmark to remove. --- tinygrad/opt/heuristic.py | 44 +++++++++++++++++++-------------------- tinygrad/opt/kernel.py | 8 +++++++ 2 files changed, 30 insertions(+), 22 deletions(-) diff --git a/tinygrad/opt/heuristic.py b/tinygrad/opt/heuristic.py index be7a990991..9146f44cf2 100644 --- a/tinygrad/opt/heuristic.py +++ b/tinygrad/opt/heuristic.py @@ -27,9 +27,9 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]: return k.applied_opts # are we grouping? (requires local shape support) - if resolve(prod(k.sts[0].shape[:k.first_reduce]) <= 2048, False): + if resolve(prod(k.sts[0].shape[i] for i in k.upcastable_dims) <= 2048, False): for sz in [16]: - try: # may fail due to excessive smem usage + try: k.apply_opt(Opt(OptOps.GROUPTOP, 0, sz)) break except KernelOptError: pass @@ -38,9 +38,9 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]: for buf_index,buf in enumerate(k.bufs): if isinstance(buf.src[0].dtype, ImageDType): if (unit_stride_axes_mul_4 := [i for i in k.sts[buf_index].unit_stride_axes(ignore_valid=True) if k.sts[buf_index].shape[i]%4 == 0]): - if k.axis_types[axis:=unit_stride_axes_mul_4[0]] in (AxisType.GLOBAL, AxisType.LOCAL): + if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims: k.apply_opt(Opt(OptOps.UPCAST, axis, 4)) - elif k.axis_types[axis:=unit_stride_axes_mul_4[0]] in (AxisType.GROUP_REDUCE, AxisType.REDUCE): + elif axis in k.unrollable_dims: k.apply_opt(Opt(OptOps.UNROLL, axis-k.first_reduce, 4)) # no more opt if we are grouping @@ -51,9 +51,8 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]: # if there are small dims with lots of valid masks, upcast them (they might be from Tensor.stack) to_upcast: list[int] = [] # upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first) - for axis in range(k.first_reduce): - # for now skip upcasting here if there is a symbolic axis - if isinstance(k.full_shape[axis], int) and k.full_shape[axis] <= 7 and any(st.axis_is_masked(axis) for st in k.sts) and \ + for axis in k.upcastable_dims: + if k.full_shape[axis] <= 7 and any(st.axis_is_masked(axis) for st in k.sts) and \ prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7: if DEBUG >= 4: print(f"upcasting masked axis : {axis}") to_upcast.append(axis) @@ -62,12 +61,12 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]: # potentially do more upcasts of non reduce axes based on a heuristic is_dsp = k.opts is not None and k.opts.device == "DSP" upcasted_axis: set[int] = set() - while resolve(prod(k.sts[0].shape[:k.first_reduce]) >= 1024): + while resolve(prod(k.sts[0].shape[i] for i in k.upcastable_dims) >= 1024): xb_choices = [] # consider all the non reduce axes, and a 3 or 4 reduce. (128 on the DSP) - for axis, upcast_amount in itertools.product(range(k.first_reduce), ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]): - # if we haven't upcasted it, it's not symbolic, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already - if axis not in upcasted_axis and isinstance(k.full_shape[axis], int) and k.full_shape[axis]%upcast_amount == 0 and \ + for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]): + # if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already + if axis not in upcasted_axis and k.full_shape[axis]%upcast_amount == 0 and \ any(st.views[-1].strides[axis] == 0 and not any(x == 0 for x in st.real_strides()[k.first_upcast:]) \ for st in k.sts): xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in k.sts), @@ -79,25 +78,25 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]: upcasted_axis.add(xb_choices[0][2]) else: break - # if last dim is small(ish) and it's a reduce dim, loop unroll the reduce - if k.first_reduce < k.first_upcast and \ + # if last reduce dim is small(ish), loop unroll the reduce + if k.unrollable_dims and \ (prod(k.full_shape[k.first_upcast:]) <= 4 or (AxisType.UNROLL not in k.axis_types)) and (prod(k.full_shape[k.first_upcast:]) < 64): - if isinstance(s:=k.full_shape[k.first_upcast-1], int) and s <= 32: # NOTE: cannot loop unroll symbolic axis - k.apply_opt(Opt(OptOps.UNROLL, k.first_upcast-1-k.first_reduce, 0)) + if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32: + k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims[-1]-k.first_reduce, 0)) # if it's small, upcast a second reduce dimension too - if k.first_reduce < k.first_upcast and s <= 3 and isinstance(s2:=k.full_shape[k.first_upcast-1], int) and s2 <= 3: - k.apply_opt(Opt(OptOps.UNROLL, k.first_upcast-1-k.first_reduce, 0)) + if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3: + k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims[-1]-k.first_reduce, 0)) else: for splits in [4]: - if k.full_shape[k.first_upcast-1]%splits == 0: - k.apply_opt(Opt(OptOps.UNROLL, k.first_upcast-1-k.first_reduce, splits)) + if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0: + k.apply_opt(Opt(OptOps.UNROLL, axis-k.first_reduce, splits)) break # if nothing at all is upcasted and it's easy to, do an upcast for splits in [4]: # TODO: somehow this never hits a reduce - if k.upcasted == 0 and k.shape_len > 0 and k.full_shape[k.first_upcast-1] % splits == 0: - k.apply_opt(Opt(OptOps.UPCAST, k.first_upcast-1, splits)) + if not k.upcasted and k.upcastable_dims and k.full_shape[k.upcastable_dims[-1]] % splits == 0: + k.apply_opt(Opt(OptOps.UPCAST, k.upcastable_dims[-1], splits)) # **** local groups **** @@ -106,7 +105,8 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]: k.apply_opt(Opt(OptOps.NOLOCALS)) else: # prioritize making expand axes local - local_axis_ranking = [(any(st.views[-1].strides[axis] == 0 for st in k.sts), axis) for axis in range(k.first_reduce)] + local_axis_ranking = [(any(st.views[-1].strides[axis] == 0 for st in k.sts), axis) \ + for axis,t in enumerate(k.axis_types) if t in (AxisType.GLOBAL, AxisType.LOOP)] to_local: list[tuple[int, int]] = [] for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])): local_size = prod(sz for _, sz in to_local) diff --git a/tinygrad/opt/kernel.py b/tinygrad/opt/kernel.py index 9611a956d9..c84d971560 100644 --- a/tinygrad/opt/kernel.py +++ b/tinygrad/opt/kernel.py @@ -141,6 +141,14 @@ class Kernel: @property def group_for_reduces(self) -> int: return sum([1 for x in self.axis_types if x == AxisType.GROUP_REDUCE]) if hasattr(self, 'axis_types') else 0 + # heuristic helpers + @property + def upcastable_dims(self) -> list[int]: return [i for i,(a,s) in enumerate(zip(self.axis_types, self.full_shape)) \ + if a in (AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) and isinstance(s, int) and s > 1] + @property + def unrollable_dims(self) -> list[int]: return [i for i,(a,s) in enumerate(zip(self.axis_types, self.full_shape)) \ + if a in (AxisType.REDUCE, AxisType.GROUP_REDUCE) and isinstance(s, int) and s > 1] + # ******************** colors and names ******************** def colors(self) -> list[str]: