Jean-Baptiste Lespiau
e95d5701e3
Add benchmarks for specifically the dispatch time. ( #4128 )
...
The goal is to distinguish the time it takes for `jitted_f` to return, and the time it takes to return and wait for the result.
We also add one to distinguish the time it takes to call the function with the argument transfer or without it.
e.g.
name time/op
jit_trivial_dispatch 28.9µs ± 2%
jit_trivial 31.5µs ± 5%
jit_simple_dispatch 60.7µs ± 4%
jit_simple 129µs ±24%
jit_simple_many_args_disptch 390µs ±19%
jit_simple_many_args 388µs ±16%
jit_dispatch_without_transfer 379µs ± 6%
jit_dispatch_with_transfer 450µs ± 5%
2020-08-27 17:02:13 +03:00
Jake Vanderplas
29aa9bfc8f
Cleanup: avoid jnp.prod & np.prod on array shapes ( #4086 )
2020-08-18 10:17:38 -07:00
Jake Vanderplas
a7c2cdea64
Cleanup: convert uses of import numpy as onp in library code ( #3754 )
2020-07-14 13:05:31 -07:00
Chris Jones
c1aeb8b3fe
Add simple JAX API microbenchmarks. ( #3679 )
2020-07-09 10:02:23 -07:00
Jake Vanderplas
6aa8f2461c
Fix remaining flakes and use exclude within setup.cfg ( #3371 )
2020-06-08 22:58:03 -07:00
Peter Hawkins
b1bc841ae5
Replace np -> jnp, onp -> np in more places. ( #2973 )
...
* Replace np -> jnp, onp -> np in more places.
Context: #2370
* Fix typo in random_test.py
2020-05-05 16:40:41 -04:00
Skye Wanderman-Milne
4b0334338e
Add pmap_shard_device_array_benchmark. ( #2864 )
...
Also renames pmap_shard_args_benchmark to pmap_shard_sharded_device_array_benchmark.
2020-04-27 17:21:05 -07:00
Skye Wanderman-Milne
8c2901cf4a
Add --export_dir and --baseline_dir flags to benchmark.py. ( #2677 )
...
`--export_dir` allows saving benchmark results to CSV files, and
`--baseline_dir` allows comparing results to a baseline exported via
`--export_dir`.
2020-04-13 10:07:05 -07:00
Skye Wanderman-Milne
3fe8bd027c
Adjust pmap_bechmark.py values to be more realistic. ( #2622 )
2020-04-06 16:38:34 -07:00
Skye Wanderman-Milne
c28c46e191
Add ShardedDeviceArray indexing benchmark. ( #2549 )
...
Example output:
```
---------Benchmark summary for ShardedDeviceArray_indexing---------
indices_fn mean %std relative
------------------ -------- ------- ----------
integer_indices 0.16901 8.52522 1
integer_2D_indices 18.4918 0 109.412
```
2020-03-31 15:52:41 -07:00
George Necula
fd52fbf411
Fix import in benchmarks
...
This works on my machine as 'python benchmarks/pmap_benchmark.py'. It also
follows the code in examples.
This will need a copybara rule to change the import to 'from jax.benchmarks import benchmark'
2020-03-31 11:48:08 +03:00
Skye Wanderman-Milne
24bbd2bc1d
Fix pmap_benchmark.py import ( #2524 )
2020-03-27 10:50:57 -07:00
George Necula
428377afb3
Added type annotations and removed unused imports ( #2472 )
...
* Added type annotations and removed unused imports
* Adjusted type hints for pytype
2020-03-21 13:54:30 +01:00
George Necula
cd7ab0a9e0
Changed to pmap_benchmark to make it runnable in Google ( #2448 )
2020-03-19 06:56:59 +01:00
Skye Wanderman-Milne
75077a1441
Add pmap_benchmark.py ( #2409 )
...
Example output:
```
$ TARGET_TOTAL_SECS=2 CUDA_VISIBLE_DEVICES= XLA_FLAGS=--xla_force_host_platform_device_count=500 python3 benchmarks/pmap_benchmark.py
2020-03-12 15:46:35.903121: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:313] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
/usr/local/google/home/skyewm/jax/jax/lib/xla_bridge.py:122: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
---------Benchmark results for pmap_shard_args_nargs=10_nshards=4---------
mean=0.034490 std=0.002890 %std=8.378140 total=2.000426
#iters=58 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=100_nshards=4---------
mean=0.091495 std=0.005935 %std=6.486871 total=2.012888
#iters=22 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=101_nshards=4---------
mean=0.113549 std=0.009080 %std=7.996712 total=2.043878
#iters=18 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=500_nshards=4---------
mean=0.356868 std=0.007960 %std=2.230518 total=2.141210
#iters=6 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=2---------
mean=0.022288 std=0.002946 %std=13.219607 total=2.005951
#iters=90 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=4---------
mean=0.035210 std=0.002024 %std=5.747389 total=2.006975
#iters=57 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=8---------
mean=0.048641 std=0.001578 %std=3.243398 total=2.042912
#iters=42 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=100---------
mean=0.257487 std=0.007190 %std=2.792452 total=2.059900
#iters=8 #warmup=1
---------Benchmark results for pmap_shard_args_nargs=10_nshards=500---------
mean=1.696294 std=0.005097 %std=0.300473 total=3.392588
#iters=2 #warmup=1
---------Benchmark summary for pmap_shard_args---------
nargs nshards mean %std relative
------- --------- --------- --------- ----------
10 4 0.0344901 8.37814 1
100 4 0.0914949 6.48687 2.65279
101 4 0.113549 7.99671 3.29221
500 4 0.356868 2.23052 10.347
10 2 0.0222883 13.2196 0.646224
10 4 0.0352101 5.74739 1.02088
10 8 0.0486408 3.2434 1.41028
10 100 0.257487 2.79245 7.46555
10 500 1.69629 0.300473 49.182
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=4---------
mean=0.061780 std=0.004737 %std=7.668032 total=2.038743
#iters=33 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=100_nshards=4---------
mean=0.123264 std=0.005980 %std=4.851385 total=2.095494
#iters=17 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=500_nshards=4---------
mean=0.471524 std=0.024051 %std=5.100792 total=2.357622
#iters=5 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=2---------
mean=0.041546 std=0.004446 %std=10.700256 total=2.035745
#iters=49 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=4---------
mean=0.063768 std=0.002756 %std=4.322039 total=2.040561
#iters=32 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=8---------
mean=0.087285 std=0.005343 %std=6.121320 total=2.007556
#iters=23 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=100---------
mean=0.623440 std=0.004038 %std=0.647725 total=2.493759
#iters=4 #warmup=1
---------Benchmark results for pmap_shard_outputs_nouts=10_nshards=500---------
mean=4.096676 std=0.000000 %std=0.000000 total=4.096676
#iters=1 #warmup=1
---------Benchmark summary for pmap_shard_outputs---------
nouts nshards mean %std relative
------- --------- --------- --------- ----------
10 4 0.0617801 7.66803 1
100 4 0.123264 4.85139 1.99521
500 4 0.471524 5.10079 7.6323
10 2 0.0415458 10.7003 0.672479
10 4 0.0637675 4.32204 1.03217
10 8 0.087285 6.12132 1.41283
10 100 0.62344 0.647725 10.0913
10 500 4.09668 0 66.3106
```
2020-03-17 14:31:25 -07:00
