mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 04:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
14,645 Commits 503 Branches 263 Tags
Commit Graph

47 Commits

Author SHA1 Message Date
Peter Hawkins
428189f8fb Replace uses of deprecated JAX sharding APIs with their new names in jax.sharding. 
This change updates:
* {jax.experimental.maps.Mesh, jax.interpreters.pxla.Mesh} to jax.sharding.Mesh
* {jax.experimental.PartitionSpec, jax.experimental.pjit.PartitionSpec, jax.interpreters.pxla.PartitionSpec, jax.pxla.PartitionSpec} to jax.sharding.PartitionSpec
* jax.experimental.maps.NamedSharding to jax.sharding.NamedSharding.

PiperOrigin-RevId: 506994892
 2023-02-03 14:28:45 -08:00
Jake VanderPlas
43e57db77a Begin deprecation of public jax.ShapedArray 2023-01-30 11:27:58 -08:00
jax authors
eb875cd5dd Added a pattern-match optimisation for inplace-select. 
PiperOrigin-RevId: 497425937
 2022-12-23 16:05:56 -08:00
Peter Hawkins
d6c67c97db Remove redundant dtype canonicalization from jax.device_put(). 
Gives a small improvement to the included jax.device_put() benchmark on my VM:

```
name        old cpu/op  new cpu/op  delta
device_put  91.3µs ± 5%  80.1µs ± 3%  -12.29%  (p=0.008 n=5+5)

name        old time/op             new time/op             delta
device_put  91.4µs ± 5%             80.1µs ± 3%  -12.29%          (p=0.008 n=5+5)
```

jax.device_put() has not been optimized that much yet and there is plenty of room for further improvement.

PiperOrigin-RevId: 491727173
 2022-11-29 13:47:36 -08:00
Yash Katariya
928dee415f Optimize host_local_array_to_global_array by caching the local to global conversion and flattening of axis resources. Also take a fast path for device_put which does not do abstractify and only canonicalize_dtype on the entire array once (instead of doing it for every shard). 
This results in a 5x speedup!

Before:

```
---------------------------------------------------------------------------
Benchmark                                 Time             CPU   Iterations
---------------------------------------------------------------------------
host_local_array_to_global_array       3.03 ms         3.02 ms          220
```

After:

```
---------------------------------------------------------------------------
Benchmark                                 Time             CPU   Iterations
---------------------------------------------------------------------------
host_local_array_to_global_array      0.673 ms        0.671 ms          985
```

PiperOrigin-RevId: 489880547
 2022-11-20 20:53:02 -08:00
Yash Katariya
c42bad85ef Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL). 
PiperOrigin-RevId: 488473538
 2022-11-14 14:44:00 -08:00
jax authors
f4be5ab173 Merge pull request #12219 from jakevdp:indexing-slice 
PiperOrigin-RevId: 485946084
 2022-11-03 12:44:28 -07:00
Yash Katariya
532cd7ed74 Skip the benchmarks properly via state.skip_with_error when enough devices are not present. 
PiperOrigin-RevId: 485931295
 2022-11-03 11:44:57 -07:00
Jake VanderPlas
753562d574 Add benchmarks for repeated static indexing & slicing 2022-11-03 11:41:37 -07:00
Hyeontaek Lim
fc8f40ce0e Internal visibility change 
PiperOrigin-RevId: 484340424
 2022-10-27 13:49:16 -07:00
Yash Katariya
cf6b5097d0 Remove pytest_benchmark for test-requirements.txt and move the benchmark file which was using that package to use google_benchmark. 
PiperOrigin-RevId: 483736267
 2022-10-25 11:59:32 -07:00
Yash Katariya
3572bb2db0 [Rollback] 
Allow uncommitted single device PyArray in C++ pjit path.

PiperOrigin-RevId: 482084898
 2022-10-18 19:42:10 -07:00
Kuangyuan Chen
d64da3d407 Roll forward with fix: Remove the original python function fun_ from C++ PjitFunction, as the destroying fun_ may yield the thread in some cases, which causes error during deleting the python object of PjitFunction. 
PiperOrigin-RevId: 481950912
 2022-10-18 10:05:53 -07:00
Kuangyuan Chen
fd2f590b3b Allow uncommitted single device PyArray in C++ pjit path. 
PiperOrigin-RevId: 481711690
 2022-10-17 12:35:30 -07:00
jax authors
504b3c1b25 roll forward with the fix: Make params arg in Compiled.call() position-only so that it does not conflict with the keyword args. 
PiperOrigin-RevId: 481666211
 2022-10-17 09:50:55 -07:00
Kuangyuan Chen
38a7582923 roll forward with the fix: Make params arg in Compiled.call() position-only so that it does not conflict with the keyword args. 
PiperOrigin-RevId: 481181330
 2022-10-14 10:42:15 -07:00
jax authors
1945208d34 Rollback because of failing tests internally. 
PiperOrigin-RevId: 481103002
 2022-10-14 03:12:42 -07:00
Kuangyuan Chen
d082ea0d46 Implement a fast path for pjit AOT in C++ for jax.Array inputs. 
PiperOrigin-RevId: 480983807
 2022-10-13 14:24:05 -07:00
Yash Katariya
9e4114f0f1 Move array.py and sharding.py from experimental/ to _src/. 
PiperOrigin-RevId: 477201711
 2022-09-27 10:06:52 -07:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". 
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
 2022-09-22 12:27:19 -07:00
Kuangyuan Chen
405a2310ce Implement pjit fast path in cpp for jax.Array inputs 
PiperOrigin-RevId: 475988677
 2022-09-21 20:18:18 -07:00
Kuangyuan Chen
b9e384363e Add many args benchmark for jax.Array 
PiperOrigin-RevId: 475853211
 2022-09-21 09:54:51 -07:00
Kuangyuan Chen
b764aadbcf Add pmap benchmarks using jax.Array 
PiperOrigin-RevId: 473325120
 2022-09-09 13:10:42 -07:00
Kuangyuan Chen
d17e516ea7 Add benchmarks for jax.Array 
PiperOrigin-RevId: 471889808
 2022-09-02 14:45:09 -07:00
Matthew Johnson
b3b4ffbc21 make slicing compilation a little faster 
For the common special case of int or slice-of-int/None indexing, we can
generate a lax.slice rather than a lax.gather. That makes compilation a little
faster, and makes the generated jaxpr a bit more wieldy too, to process in
transformations and to read when pretty-printed.
 2022-08-19 09:14:37 -07:00
Matthew Johnson
42dd7cac43 simplify slicing jaxprs a little 
Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com>
 2022-08-12 19:52:02 -07:00
Matthew Johnson
e3a92d52ba prepare to switch to new remat 
This commit involves a few things, which are all united in being about landing
the new remat (aka new checkpoint) implementation:
  * add benchmarks for new remat eager performance, and some caching to make those
    benchmarks fast
  * warn when the old-remat-exclusive `concrete` feature is used, with an
    actionable message pointing to the new recommended approach involving static_argnums
  * add the static_argnums parameter to both new and old remt
  * update docstrings (and de-duplicate them to)
  * add new tests, especially around caching and errors/warnings
 2022-08-04 12:25:03 -07:00
Yash Katariya
2109c6ec8c Make is_compatible_aval an optional method which sharding interfaces can implement to raise a more meaningful error. Otherwise lower to opsharding and catch the error if it fails. 
PiperOrigin-RevId: 464147877
 2022-07-29 13:37:15 -07:00
Yash Katariya
47623264db Export HloSharding via pybind which is a C++ wrapper around OpSharding proto. 
PiperOrigin-RevId: 463992136
 2022-07-28 21:01:15 -07:00
Matthew Johnson
148173630f add an optional fastpath for api_util.shaped_abstractify 
also add a benchmark for it, 8.7ms -> 0.2ms on my machine

Co-authored-by: Yash Katariya <yashkatariya@google.com>
 2022-07-27 15:14:37 -07:00
Jake VanderPlas
a10f0377db Avoid top-level aliases of jax.tree_util.* 2022-07-07 11:41:02 -07:00
Qiao Zhang
5d7f639769 Add small and big matmul to api_benchmarks. 
name                                  cpu/op
jit_small_matmul                      2.96µs ± 2%
jit_big_matmul                        22.1µs ±21%

name                                  time/op

jit_small_matmul                      2.96µs ± 2%
jit_big_matmul                        22.7µs ±21%

PiperOrigin-RevId: 435453853
 2022-03-17 14:51:17 -07:00
Tianjian Lu
6a0bfdd637 [JAX] Requires indices to be sorted and of int32 in _sparse_bcoo_matvec test. 
PiperOrigin-RevId: 417695937
 2021-12-21 14:56:52 -08:00
Jake VanderPlas
b8296bb06c benchmarks: add compilation of sparse operations 
PiperOrigin-RevId: 405752775
 2021-10-26 15:41:21 -07:00
Jake VanderPlas
c62452f2d2 benchmarks: add JIT versions of sparse.BCOO benchmarks 
PiperOrigin-RevId: 405696495
 2021-10-26 11:39:01 -07:00
Jake VanderPlas
5a3733788c benchmarks: add sparse.BCOO todense/fromdense/matvec benchmarks 
PiperOrigin-RevId: 405519824
 2021-10-25 16:44:36 -07:00
Yash Katariya
bfbdfa87e7 Add a warmup loop to pmap_simple_8_devices_100_args benchmark so as to not measure the compile time. 
PiperOrigin-RevId: 401402336
 2021-10-06 19:51:35 -07:00
Peter Hawkins
6a1b626564 Remove jax.api. 
Functions exported as jax.api were aliases for names in jax.*. Use the jax.* names instead.
 2021-09-16 16:29:06 -04:00
Jean-Baptiste Lespiau
6cb8737c1a Add a benchmark with many arguments. 
PiperOrigin-RevId: 393216026
 2021-08-26 15:06:09 -07:00
Jean-Baptiste Lespiau
afdd195e42 Internal only. 
PiperOrigin-RevId: 388562206
 2021-08-03 15:49:17 -07:00
Qiao Zhang
850bd66242 [JAX] Prune unused inputs in jit. 
- Python part based on: https://github.com/google/jax/pull/6567
- Added cpp_jit path to handle pruned args

PiperOrigin-RevId: 371743277
 2021-05-03 11:41:29 -07:00
Tom Hennigan
9d56552517 Remove special casing on npy_value when indexing sharded arrays. 
Before:

```
GPU
sda_index_2                           2912972 ns      2778716 ns          256

TPU
sda_index_1                            769968 ns       751700 ns          921
sda_index_2                           1510841 ns      1489716 ns          465
sda_index_8                           6102259 ns      6027655 ns          117
```

After:

```
GPU
sda_index_2                             28095 ns        27983 ns        25463

TPU
sda_index_1                             10302 ns        10279 ns        67884
sda_index_2                             20010 ns        19947 ns        34628
sda_index_8                             78492 ns        78306 ns         8934
```

PiperOrigin-RevId: 368380864
 2021-04-14 01:17:10 -07:00
Matthew Johnson
9802f3378e add simple single-primitive eager benchmarks 2021-03-18 21:46:46 -07:00
Peter Hawkins
cdd36b1113 Improve API benchmarks. 
Add benchmarks for different dispatch arg arities.

Add more blocking before and after benchmark loops that don't otherwise block.
 2021-03-03 20:50:45 -05:00
Tom Hennigan
7b94c44af9 Sort imports. 
PiperOrigin-RevId: 359532656
 2021-02-25 08:52:09 -08:00
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
Chris Jones
c1aeb8b3fe
Add simple JAX API microbenchmarks. (#3679) 2020-07-09 10:02:23 -07:00