mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 11:16:07 +00:00
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.
This commit is contained in:
chenyu
+22
-22
@@ -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)
|
||||
|
||||
@@ -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]:
|
||||
|
||||
Reference in New Issue
Block a user