Update251203 (#233)
This commit is contained in:
carrot
@@ -1,10 +1,50 @@
|
||||
import itertools
|
||||
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError, AxisType
|
||||
from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS
|
||||
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
|
||||
from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, AMX
|
||||
from tinygrad.dtype import ImageDType
|
||||
from tinygrad.uop.ops import Ops, resolve
|
||||
from tinygrad.uop.ops import Ops, resolve, AxisType
|
||||
from tinygrad.codegen.opt.postrange import Scheduler
|
||||
|
||||
def hand_coded_optimizations(k:Scheduler) -> Scheduler:
|
||||
# first try the tensor cores
|
||||
""" Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
|
||||
Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
|
||||
|
||||
Keyword arguments:
|
||||
use_tensor_cores -- controls how tensor cores are applied (default 1)
|
||||
0: will disable any tensor core matching
|
||||
1: enable tensor cores
|
||||
2: apply tensor core shape but don't use UOp.WMMA
|
||||
extra_opts -- additional Opt's to apply after the tensor core instead of the hand-coded additional Opt's (default None)
|
||||
tc_select -- specifies which tensor core(s) to use for optimization (default -1)
|
||||
-1: iterates through all available tensor cores in order and uses the first one that matches the requirements (dims and dtypes)
|
||||
[0-N]: uses only the n'th tensor core available; useful for search
|
||||
tc_opt -- controls which kinds of kernels may be eligible for tensor cores application (default 2 during BEAM, 0 otherwise)
|
||||
0: applies to only kernels with a single reduce axis and direct Ops.LOAD into Ops.MUL
|
||||
1: allows kernels with multiple reduce axes and also multiplication of Ops.CAST'd buffers
|
||||
2: allows kernels with M, N, K axes that are not multiples of the tensor core dimensions by applying padding those axes as needed
|
||||
"""
|
||||
# NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis
|
||||
if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
|
||||
good_tc_opt = False
|
||||
try: # check TC first and apply hand-coded opts if successful
|
||||
tk = k.copy()
|
||||
rngs = tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
|
||||
good_tc_opt = True
|
||||
except KernelOptError:
|
||||
pass
|
||||
if good_tc_opt:
|
||||
# skip hand-coded TC opts if AMX, upcasting will make kernel slower
|
||||
if rngs is not None and not AMX:
|
||||
for tc_dim in [1,0]: # attempt to upcast M and N
|
||||
szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
|
||||
if szs:
|
||||
# set it to the replaced range
|
||||
rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
|
||||
if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
|
||||
tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
|
||||
return tk
|
||||
|
||||
def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
# make a copy so it does not mutate the input
|
||||
k = k.copy()
|
||||
|
||||
@@ -13,19 +53,17 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
if k.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and \
|
||||
k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.opts.has_shared and \
|
||||
(mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
|
||||
st0, st1 = k.sts[k.bufs.index(mulop.src[0])], k.sts[k.bufs.index(mulop.src[1])]
|
||||
strides0, strides1 = st0.real_strides(), st1.real_strides()
|
||||
def has_expanded_axis(shape, strides): return any(resolve(s > 1) and not resolve(st != 0) for s,st in zip(shape,strides))
|
||||
if strides0[first_reduce:=(k.axes_of(AxisType.REDUCE)[0])] == 1 and \
|
||||
not (has_expanded_axis(st0.shape, strides0) and has_expanded_axis(st1.shape, strides1)):
|
||||
idx0, idx1 = mulop.src[0].src[0].src[1].get_idx(), mulop.src[1].src[0].src[1].get_idx()
|
||||
first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
|
||||
if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
|
||||
for global_idx in k.axes_of(AxisType.GLOBAL):
|
||||
if k.full_shape[first_reduce]%MV_THREADS_PER_ROW == 0 and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
|
||||
if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
|
||||
if DEBUG >= 3:
|
||||
print(f"MATVEC: {k.full_shape=} {first_reduce=} {strides0=} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
|
||||
print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
|
||||
if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
|
||||
if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
|
||||
if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
|
||||
return k.applied_opts
|
||||
return k
|
||||
|
||||
# are we grouping? (requires local shape support)
|
||||
if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
|
||||
@@ -38,14 +76,17 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
# upcast float4 images
|
||||
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]):
|
||||
# part of real_strides
|
||||
unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].get_idx().split_uop(Ops.ADD) if
|
||||
c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
|
||||
if len(unit_stride_axes_mul_4):
|
||||
if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims:
|
||||
k.apply_opt(Opt(OptOps.UPCAST, axis, 4))
|
||||
elif axis in k.unrollable_dims:
|
||||
k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims.index(axis), 4))
|
||||
|
||||
# no more opt if we are grouping
|
||||
if k.group_for_reduces: return k.applied_opts
|
||||
if k.group_for_reduces: return k
|
||||
|
||||
# **** below this line need to be optional and benchmarked ****
|
||||
|
||||
@@ -53,8 +94,9 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
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 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:
|
||||
# for Schedule, we check if the range is used in INDEX gates or WHERE gates
|
||||
is_masked = any(any(o is k.rngs[axis] for o in u.src[0].parents) for u in k.ast.parents if u.op is Ops.WHERE)
|
||||
if k.full_shape[axis] <= 7 and is_masked 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)
|
||||
for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0))
|
||||
@@ -68,10 +110,18 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
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 in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
|
||||
if any(st.views[-1].strides[axis] == 0 and \
|
||||
all(x != 0 for t,x in zip(k.axis_types, st.real_strides()) if t in (AxisType.UPCAST, AxisType.UNROLL)) for st in k.sts):
|
||||
xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in k.sts),
|
||||
sum(st.views[-1].strides[axis] for st in k.sts), axis, upcast_amount))
|
||||
rng = k.rngs[axis]
|
||||
if any(rng not in b.src[1].get_idx().parents and all(r2 in b.src[1].get_idx().parents
|
||||
for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs):
|
||||
num_strides, sum_strides = 0, 0
|
||||
for b in k.bufs:
|
||||
idx = b.src[1].get_idx()
|
||||
if rng in idx.parents: num_strides += 1
|
||||
for c in idx.split_uop(Ops.ADD):
|
||||
if c is rng: sum_strides += 1
|
||||
if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
|
||||
if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
|
||||
xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
|
||||
if xb_choices:
|
||||
xb_choices = sorted(xb_choices)
|
||||
if DEBUG >= 4: print(f"more upcast axis : {xb_choices}")
|
||||
@@ -109,7 +159,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 k.axes_of(AxisType.GLOBAL, AxisType.LOOP)]
|
||||
local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().parents for b in k.bufs), axis) \
|
||||
for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
|
||||
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)
|
||||
@@ -122,4 +173,16 @@ def hand_coded_optimizations(k:Kernel) -> list[Opt]:
|
||||
k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
|
||||
if will_delete_shape: deleted_shape += 1
|
||||
|
||||
return k.applied_opts
|
||||
# **** threading ****
|
||||
|
||||
if k.opts.has_threads and k.opts.global_max is not None:
|
||||
for threads in [32,16,12,8,6,5,4,3,2]:
|
||||
# Skip is too many threads. Heuristic: use about 128K ops per thread
|
||||
if threads > k.opts.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
|
||||
for axis in k.axes_of(AxisType.LOOP):
|
||||
if k.full_shape[axis] % threads == 0:
|
||||
k.apply_opt(Opt(OptOps.THREAD, axis, threads))
|
||||
break
|
||||
if k.applied_opts and k.applied_opts[-1].op is OptOps.THREAD: break
|
||||
|
||||
return k
|
||||
|
||||
Reference in New Issue
Block a user