rocm_jax/benchmarks
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
..