github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-11 07:58:08 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add Opt(op=OptOps.LOCAL, axis=6, amt=2) to actions (#7147)

Browse Source 
* 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 <module>
    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
This commit is contained in:
chenyu
2024-10-18 09:03:24 -04:00
committed by GitHub
parent 4cf7cca91a
commit ea2efbf508
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tinygrad/engine/search.py
View File

@@ -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)]