OkanoShinri
commented
11 months ago benchmark.py
```py from pgx.mahjong._env import ( Mahjong, _discard, _selfkan, _riichi, _tsumo, _ron, _pon, _minkan, _pass, ) import jax import time import sys import timeit # func(state, action) functions1 = {"_discard": _discard, "_selfkan": _selfkan} # func(state) functions2 = { "_riichi": _riichi, "_tsumo": _tsumo, "_ron": _ron, "_pon": _pon, "_minkan": _minkan, "_pass": _pass, } env = Mahjong() N = 10 func_name = sys.argv[1] if func_name in functions1: func = functions1[func_name] key = jax.random.PRNGKey(352) state = env.init(key=key) time_sta = time.perf_counter() jax.jit(func)(state, 0) time_end = time.perf_counter() delta = (time_end - time_sta) * 1000 exp = jax.make_jaxpr(func)(state, 0) n_line = len(str(exp).split("\n")) jit_func = jax.jit(func) run_delta = timeit.timeit( "jit_func(state, 0)", globals=globals(), number=N ) print( f"| `{func.__name__}` | {n_line} | {delta:.1f}ms | {run_delta/N*1000000:.1f}μs |" ) elif func_name in functions2: func = functions2[func_name] key = jax.random.PRNGKey(352) state = env.init(key=key) time_sta = time.perf_counter() jax.jit(func)(state) time_end = time.perf_counter() delta = (time_end - time_sta) * 1000 exp = jax.make_jaxpr(func)(state) n_line = len(str(exp).split("\n")) jit_func = jax.jit(func) run_delta = timeit.timeit("jit_func(state)", globals=globals(), number=N) print( f"| `{func.__name__}` | {n_line} | {delta:.1f}ms | {run_delta/N*1000000:.1f}μs |" ) ```benchmark.sh
```py echo "| function name | # expr lines | compile time | running time |" echo "| :--- | ---: | ---: | ---: |" for funcname in _discard _selfkan _riichi _tsumo _ron _pon _minkan _pass do python3 benchmark.py $funcname done ```- 手元での実行時間
| function name | # expr lines | compile time | running time |
|---|---|---|---|
_discard |
12245 | 5012.2ms | 264.3μs |
- _draw |
5517 | 2137.6ms | 244.9μs |
-- _make_legal_action_mask |
3048 | 1281.3ms | 106.2μs |
-- _make_legal_action_mask_w_riichi |
2394 | 1088.4ms | 78.7μs |
_selfkan |
333 | 198.3ms | 64.7μs |
_riichi |
934 | 427.5ms | 104.6μs |
_tsumo |
2328 | 1095.3ms | 77.1μs |
_ron |
2280 | 1048.6ms | 104.6μs |
_pon |
261 | 117.2ms | 64.7μs |
_minkan |
324 | 133.0ms | 69.7μs |
_pass |
8539 | 2526.0ms | 223.7μs |
can_tsumo |
225 | 115.3ms | 12.5μs |
can_ron |
245 | 122.4ms | 10.6μs |
can_minkan |
9 | 12.9ms | 8.6μs |
can_chi |
206 | 63.3ms | 11.7μs |
can_riichi |
424 | 289.8ms | 121.8μs |
is_tenpai |
344 | 189.0ms | 11.9μs |
- 優先順位
- [ ]
discard - [ ]
pass - [ ]
ron - [ ]
tsumo
sotetsuk
benchmark.py
```py from pgx._mahjong._mahjong2 import ( Mahjong, _discard, _selfkan, _riichi, _tsumo, _ron, _pon, _minkan, _pass, ) import jax import time import sys # func(state, action) functions1 = {"_discard": _discard, "_selfkan": _selfkan} # func(state) functions2 = { "_riichi": _riichi, "_tsumo": _tsumo, "_ron": _ron, "_pon": _pon, "_minkan": _minkan, "_pass": _pass, } env = Mahjong() func_name = sys.argv[1] if func_name in functions1: func = functions1[func_name] key = jax.random.PRNGKey(352) state = env.init(key=key) time_sta = time.perf_counter() jax.jit(func)(state, 0) time_end = time.perf_counter() delta = (time_end - time_sta) * 1000 exp = jax.make_jaxpr(func)(state, 0) n_line = len(str(exp).split("\n")) print(f"| `{func.__name__}` | {n_line} | {delta:.1f}ms |") elif func_name in functions2: func = functions2[func_name] key = jax.random.PRNGKey(352) state = env.init(key=key) time_sta = time.perf_counter() jax.jit(func)(state) time_end = time.perf_counter() delta = (time_end - time_sta) * 1000 exp = jax.make_jaxpr(func)(state) n_line = len(str(exp).split("\n")) print(f"| `{func.__name__}` | {n_line} | {delta:.1f}ms |") ```benchmark.sh
```py echo "| function name | # expr lines | compile time |" echo "| :--- | ---: | ---: |" for funcname in _discard _selfkan _riichi _tsumo _ron _pon _minkan _pass do python3 benchmark.py $funcname done ```step内の主要関数抜粋_discard_selfkan_riichi_tsumo_ron_pon_minkan_pass