From ea2efbf508b005267de388bd7fad3f2a63672ce2 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 18 Oct 2024 09:03:24 -0400 Subject: [PATCH] Add Opt(op=OptOps.LOCAL, axis=6, amt=2) to actions (#7147) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add Opt(op=OptOps.LOCAL, axis=6, amt=2) to actions it's missing if we rebuild all kernels, not just the first 2k. ``` PYTHONPATH="." GPU=1 python3 extra/optimization/get_action_space.py 29%|█████████████████████████████████████▋ | 3682/12701 [01:42<04:11, 35.83it/s]Traceback (most recent call last): File "/Users/chenyu/code/tinygrad/extra/optimization/get_action_space.py", line 27, in test_rebuild(lin) File "/Users/chenyu/code/tinygrad/extra/optimization/get_action_space.py", line 11, in test_rebuild assert o in actions, f"{o} is not in actions" ^^^^^^^^^^^^ AssertionError: Opt(op=OptOps.LOCAL, axis=6, amt=2) is not in actions ``` * break --- tinygrad/engine/search.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tinygrad/engine/search.py b/tinygrad/engine/search.py index c5061e3106..e4172f0e4f 100644 --- a/tinygrad/engine/search.py +++ b/tinygrad/engine/search.py @@ -17,7 +17,8 @@ actions += [Opt(op=OptOps.LOCAL, axis=axis, amt=amt) for amt in [2,3,4,8,13,16,2 actions += [Opt(op=OptOps.GROUPTOP, axis=axis, amt=amt) for amt in [13,16,28,29,32,49,64,256] for axis in range(3)] actions += [Opt(op=OptOps.GROUP, axis=axis, amt=amt) for amt in [0,4,8,16] for axis in range(3)] if getenv("BEAM_PADTO", 1): actions += [Opt(op=OptOps.PADTO, axis=axis, amt=amt) for amt in [32] for axis in range(7)] -actions += [Opt(op=OptOps.LOCAL, axis=0, amt=32), Opt(op=OptOps.UPCASTMID, axis=1, amt=4), Opt(op=OptOps.TC, axis=0, amt=0)] +actions += [Opt(op=OptOps.LOCAL, axis=0, amt=32), Opt(op=OptOps.LOCAL, axis=6, amt=2)] +actions += [Opt(op=OptOps.UPCASTMID, axis=1, amt=4), Opt(op=OptOps.TC, axis=0, amt=0)] actions += [Opt(op=OptOps.TC, axis=axis, amt=getenv("TC_OPT", 2)) for axis in range(9)] # covers resnet kernels (3 global * 3 reduce) actions += [Opt(op=OptOps.SWAP, axis=axis, amt=amt) for axis in range(5) for amt in range(axis+1, 5)] if getenv("NOLOCALS"): actions += [Opt(op=OptOps.NOLOCALS)]