forked from tinygrad/tinygrad
make POSTOPT=2 the default (#12034)
* make POSTOPT=2 the default * more matching tc * fix winograd * fix that test * add matvec to Scheduler * flip tc sort order * similar speed * fix beam on image * disable slow tests * slow
This commit is contained in:
George Hotz
@@ -122,8 +122,8 @@ jobs:
|
||||
run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
|
||||
- name: UsbGPU copy speeds
|
||||
run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
|
||||
- name: UsbGPU openpilot test
|
||||
run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
|
||||
#- name: UsbGPU openpilot test
|
||||
# run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (Mac)
|
||||
@@ -319,10 +319,10 @@ jobs:
|
||||
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
|
||||
- name: Run MLPerf resnet eval on training data
|
||||
run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
|
||||
#- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
# run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
#- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
# run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
|
||||
- name: Run 10 MLPerf Bert training steps (6 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
|
||||
@@ -517,10 +517,10 @@ jobs:
|
||||
run: BENCHMARK_LOG=cifar_10steps_half_wino AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
|
||||
- name: Run full CIFAR training steps w 6 GPUS
|
||||
run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
|
||||
- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
|
||||
run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
|
||||
#- name: Run full CIFAR training steps w 6 GPUS
|
||||
# run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
|
||||
#- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
|
||||
# run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: Speed (AMD Training)
|
||||
@@ -570,10 +570,10 @@ jobs:
|
||||
run: test/external/process_replay/reset.py
|
||||
- name: Run MLPerf resnet eval
|
||||
run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
|
||||
#- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
# run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
|
||||
#- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
|
||||
# run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
|
||||
- name: Run 10 MLPerf Bert training steps (6 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
|
||||
@@ -759,8 +759,8 @@ jobs:
|
||||
run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
|
||||
- name: Run full CIFAR training w 1 GPU
|
||||
run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
|
||||
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
|
||||
#- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
|
||||
# run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
|
||||
- name: Run 10 MLPerf Bert training steps (1 gpu)
|
||||
# TODO: remove BERT_LAYERS once scheduler is fast
|
||||
run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
|
||||
|
||||
@@ -468,7 +468,7 @@ jobs:
|
||||
llvm: 'true'
|
||||
- name: Test openpilot model kernel count and gate usage
|
||||
run: |
|
||||
PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2134 ALLOWED_GATED_READ_IMAGE=13 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
|
||||
PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2175 ALLOWED_GATED_READ_IMAGE=16 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
|
||||
- name: Test openpilot alt model correctness (float32)
|
||||
run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
|
||||
- name: Test openpilot fastvits model correctness (float32)
|
||||
|
||||
@@ -32,27 +32,6 @@ class TestLinearizerFailure(unittest.TestCase):
|
||||
_ = get_program(ast, Device["METAL"].renderer)
|
||||
|
||||
class TestLinearizerDumb(unittest.TestCase):
|
||||
@unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
|
||||
def test_unmerged_ifs(self):
|
||||
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(1605632), arg=0, src=())
|
||||
c1 = c0.view(ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 1, 1, 1), strides=(25088, 0, 49, 7, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)))
|
||||
c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(1605632), arg=1, src=())
|
||||
c3 = c2.view(ShapeTracker(views=(View(shape=(1, 64, 1, 512, 4, 9, 4, 9), strides=(0, 25088, 0, 49, 0, 7, 0, 1), offset=-8, mask=((0, 1), (0, 64), (0, 1), (0, 512), (0, 4), (1, 8), (0, 4), (1, 8)), contiguous=False), View(shape=(64, 1, 512, 7, 7, 512, 3, 3), strides=(663552, 0, 0, 36, 1, 1296, 360, 10), offset=0, mask=None, contiguous=False))))
|
||||
c4 = c3.load()
|
||||
c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(2359296), arg=2, src=())
|
||||
c6 = c5.view(ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 512, 3, 3), strides=(0, 0, 4608, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)))
|
||||
c7 = c6.load()
|
||||
c8 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(64, 1, 512, 7, 7, 1, 1, 1), strides=(0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=())
|
||||
c9 = c1.store(((c4*c7).cast(dtypes.float).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (5, 6, 7))).cast(dtypes.half)*UOp.const(dtypes.half, 0.9999950000374996, src=c8)).alu(Ops.MAX, UOp.const(dtypes.half, 0.0, src=c8)))
|
||||
ast = c9.sink()
|
||||
opts = [Opt(op=OptOps.TC, axis=2, arg=(-1, 2, 1)), Opt(op=OptOps.UPCAST, axis=2, arg=0), Opt(op=OptOps.UNROLL, axis=1, arg=0)]
|
||||
prg = get_program(ast, Device["METAL"].renderer, opts)
|
||||
print(prg.src)
|
||||
Device[Device.DEFAULT].compiler.compile_cached(prg.src)
|
||||
gate_count = len([x for x in prg.src.splitlines() if "if" in x])
|
||||
assert gate_count == 1, f"must have only one gate {gate_count} != 1"
|
||||
assert len([u for u in prg.uops if u.op is Ops.IF]) == 1, "must have a single IF"
|
||||
|
||||
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
|
||||
def test_max_simplify_and_cancel(self):
|
||||
c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1000), arg=0, src=())
|
||||
|
||||
@@ -3,7 +3,7 @@ import unittest, pytest
|
||||
from tinygrad import dtypes, Variable
|
||||
from tinygrad.dtype import AddrSpace
|
||||
from tinygrad.helpers import DEBUG, Context
|
||||
from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp
|
||||
from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, KernelInfo
|
||||
from tinygrad.uop.symbolic import sym
|
||||
from tinygrad.codegen import full_rewrite, full_rewrite_to_sink
|
||||
from tinygrad.codegen.late.expander import expander
|
||||
@@ -17,7 +17,7 @@ simple_pm = PatternMatcher([
|
||||
|
||||
def to_uops_list(u:List[UOp]) -> List[UOp]:
|
||||
# we strip the SINK here for legacy reasons
|
||||
ret = full_rewrite(UOp.sink(*u))
|
||||
ret = full_rewrite(UOp.sink(*u, arg=KernelInfo(opts_to_apply=())))
|
||||
assert ret[-1].op is Ops.SINK
|
||||
return ret[:-1]
|
||||
|
||||
@@ -445,7 +445,7 @@ class TestUOpGraph(unittest.TestCase):
|
||||
with Context(IGNORE_OOB=0):
|
||||
glbl0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(16), src=(), arg=0)
|
||||
v = Variable("v", 0, 20)
|
||||
st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v), UOp.const(dtypes.int, 0), UOp(Ops.IF, src=(v<16,))))
|
||||
st0 = UOp(Ops.STORE, dtypes.void, src=(glbl0.index(v, v<16), UOp.const(dtypes.int, 0)))
|
||||
to_uops_list([st0])
|
||||
|
||||
st1 = UOp(Ops.STORE, dtypes.void, (glbl0.index(v), v, v<20))
|
||||
@@ -463,7 +463,7 @@ class TestUOpGraph(unittest.TestCase):
|
||||
|
||||
gate = (gidx<400) & (lidx<8)
|
||||
|
||||
local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx), UOp.const(dtypes.uint, 1), UOp(Ops.IF, src=(lidx<8,))))
|
||||
local_store = UOp(Ops.STORE, dtypes.void, (sbuf.index(lidx, lidx<8), UOp.const(dtypes.uint, 1)))
|
||||
|
||||
barrier = UOp(Ops.BARRIER, dtypes.void, (local_store,))
|
||||
if_barrier = UOp(Ops.IF, dtypes.void, (gate, barrier))
|
||||
|
||||
@@ -29,10 +29,14 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
|
||||
"""
|
||||
# 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 isinstance(k, Kernel) and (tc_opts:=tk.tensor_core_opts) is not None and not AMX:
|
||||
# hand-coded TC opts
|
||||
@@ -43,19 +47,14 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
|
||||
if tc_opts.axes_exist[0] and (szs := [sz for sz in [4,2] if tk.full_shape[tc_opts.axes[0]] % sz == 0]): # attempt to local N
|
||||
tk.apply_opt(Opt(OptOps.LOCAL, tc_opts.axes[0], szs[0]))
|
||||
elif isinstance(k, Scheduler) and rngs is not None and not AMX:
|
||||
axes = [tk.rngs.index(r) if r in tk.rngs else None for r in rngs]
|
||||
for tc_dim in [1,0]: # attempt to upcast M and N
|
||||
if axes[tc_dim] is None: continue
|
||||
szs = [sz for sz in [5,4,3,2] if tk.full_shape[axes[tc_dim]] % sz == 0]
|
||||
szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
|
||||
if szs:
|
||||
axis = axes[tc_dim]
|
||||
if tk.full_shape[axis] == szs[0]: axes[tc_dim] = None
|
||||
tk.apply_opt(Opt(OptOps.UPCAST, axis, szs[0]))
|
||||
if axes[0] is not None and (szs := [sz for sz in [4,2] if tk.full_shape[axes[0]] % sz == 0]): # attempt to local N
|
||||
tk.apply_opt(Opt(OptOps.LOCAL, axes[0], szs[0]))
|
||||
# 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.applied_opts
|
||||
except KernelOptError:
|
||||
pass
|
||||
|
||||
# make a copy so it does not mutate the input
|
||||
k = k.copy()
|
||||
@@ -79,6 +78,18 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
|
||||
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
|
||||
else:
|
||||
idx0, idx1 = mulop.src[0].src[0].src[1], mulop.src[1].src[0].src[1]
|
||||
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 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_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
|
||||
|
||||
# are we grouping? (requires local shape support)
|
||||
if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= 2048, False):
|
||||
@@ -112,7 +123,10 @@ def hand_coded_optimizations(k:Kernel|Scheduler) -> list[Opt]:
|
||||
# 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 isinstance(k, Kernel): is_masked = any(st.axis_is_masked(axis) for st in k.sts)
|
||||
else: is_masked = any(len(st.src) > 2 and k.rngs[axis] in st.src[2].parents for st in k.bufs)
|
||||
else:
|
||||
# for Schedule, we check if the range is used in INDEX gates or WHERE gates
|
||||
is_masked = any(len(st.src) > 2 and k.rngs[axis] in st.src[2].parents for st in k.bufs) or \
|
||||
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)
|
||||
|
||||
@@ -5,7 +5,7 @@ from typing import cast, Final, Sequence
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo, graph_rewrite, _substitute, AxisType, ssimplify, can_pad
|
||||
from tinygrad.uop.symbolic import symbolic_flat
|
||||
from tinygrad.device import Buffer
|
||||
from tinygrad.dtype import AddrSpace, dtypes
|
||||
from tinygrad.dtype import AddrSpace, dtypes, ImageDType
|
||||
from tinygrad.helpers import colored, BEAM, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, POSTOPT, prod
|
||||
from tinygrad.codegen.opt import axis_colors, Opt, OptOps, KernelOptError, check, axis_letters
|
||||
from tinygrad.renderer import Renderer
|
||||
@@ -167,6 +167,7 @@ class Scheduler:
|
||||
OptOps.UNROLL: AxisType.UNROLL, OptOps.GROUP: AxisType.GROUP_REDUCE,
|
||||
OptOps.GROUPTOP: AxisType.GROUP_REDUCE}
|
||||
|
||||
ret = None
|
||||
if opt.op in opt_to_at:
|
||||
amt:int = int(rng.vmax+1) if opt.arg == 0 else cast(int, opt.arg)
|
||||
|
||||
@@ -182,7 +183,7 @@ class Scheduler:
|
||||
check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
|
||||
if opt.op is OptOps.UPCAST:
|
||||
check((self.opts is not None and self.opts.device == "DSP") or amt <= 16, "don't upcast more than 16")
|
||||
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, "upcast is for GLOBAL/LOCAL/LOOP")
|
||||
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
|
||||
if opt.op is OptOps.LOCAL:
|
||||
check(not self.dont_use_locals, "can't use locals")
|
||||
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
|
||||
@@ -190,7 +191,7 @@ class Scheduler:
|
||||
check(all(x.op is not OptOps.TC for x in self.applied_opts), "no grouping with tensor cores") # TODO: why is this wrong?
|
||||
check(not self.dont_use_locals, "can't use locals")
|
||||
check(rng.arg[-1] == AxisType.REDUCE, "group is for reduce")
|
||||
self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op==OptOps.GROUPTOP)
|
||||
ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op==OptOps.GROUPTOP)
|
||||
elif opt.op is OptOps.TC:
|
||||
check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
|
||||
check(opt.axis is not None, "tensor core opts must have an axis")
|
||||
@@ -200,8 +201,6 @@ class Scheduler:
|
||||
check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
|
||||
ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt)
|
||||
check(ret is not None, "no tensor core available")
|
||||
if append_opt: self.applied_opts.append(opt)
|
||||
return ret
|
||||
elif opt.op is OptOps.PADTO:
|
||||
check(rng.src[0].op is Ops.CONST, "only pad const axes")
|
||||
check(rng.arg[-1] not in {AxisType.UPCAST, AxisType.UNROLL}, "cannot pad upcasted") # TODO: why is this wrong?
|
||||
@@ -231,6 +230,7 @@ class Scheduler:
|
||||
raise KernelOptError(f"unsupported opt {opt.op}")
|
||||
|
||||
if append_opt: self.applied_opts.append(opt)
|
||||
return ret
|
||||
|
||||
def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
|
||||
reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE]
|
||||
@@ -247,9 +247,9 @@ class Scheduler:
|
||||
for tc in tensor_cores:
|
||||
if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
|
||||
# tensor cores have three ranges. X, Y, and REDUCE
|
||||
in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0])
|
||||
in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0])
|
||||
red_ranges = sorted(reduceop.src[1:], key=lambda x: x.arg[0])
|
||||
in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: -x.arg[0])
|
||||
in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: -x.arg[0])
|
||||
red_ranges = sorted(reduceop.src[1:], key=lambda x: -x.arg[0])
|
||||
if DEBUG >= 3:
|
||||
print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
|
||||
f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")
|
||||
@@ -341,7 +341,7 @@ class Scheduler:
|
||||
|
||||
def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
|
||||
glbls = sorted([x for x in ast.parents if x.op is Ops.DEFINE_GLOBAL], key=lambda x: x.arg)
|
||||
return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls]
|
||||
return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]
|
||||
|
||||
def apply_opts(ctx:Renderer, ast:UOp):
|
||||
if ast.tag is not None: return None
|
||||
|
||||
+1
-1
@@ -140,7 +140,7 @@ DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES = ContextVar("DONT_REALIZE_EXPAND", 0),
|
||||
QUANTIZE, VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("QUANTIZE", 0), ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
|
||||
CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
|
||||
ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, AMD_LLVM = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0), ContextVar("AMD_LLVM", 1)
|
||||
RANGEIFY, POSTOPT, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("POSTOPT", 0), ContextVar("FUSE_ATTENTION", 0)
|
||||
RANGEIFY, POSTOPT, FUSE_ATTENTION = ContextVar("RANGEIFY", 0), ContextVar("POSTOPT", 2), ContextVar("FUSE_ATTENTION", 0)
|
||||
EMULATE = ContextVar("EMULATE", "")
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
||||
Reference in New Issue
Block a user