mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-21 22:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
20,631 Commits 508 Branches 263 Tags
Commit Graph

78 Commits

Author SHA1 Message Date
jax authors
fb65ba4adf Add a config for using Clang on Windows. 
PiperOrigin-RevId: 631112031
 2024-05-06 10:39:28 -07:00
Adam Paszke
8e3f5b1018 Initial commit for Mosaic GPU 
Moving this to JAX to make it easier to explore Pallas integration.

PiperOrigin-RevId: 625982382
 2024-04-18 04:04:10 -07:00
jax authors
0be07e6aec Remove support for CUDA 11. 
Pin minimal required versions for CUDA to 12.1.

Reverts 910a31d7b7510e3375718ab1ea0d38df7bd2c0d5

PiperOrigin-RevId: 618911489
 2024-03-25 11:46:39 -07:00
jax authors
910a31d7b7 Reverts bed4f65438a62777ed100ecec2b0eb3f7cf87a0e 
PiperOrigin-RevId: 618249855
 2024-03-22 12:10:53 -07:00
jax authors
bed4f65438 Remove support for CUDA 11. 
Pin minimal required versions for CUDA to 12.1.

PiperOrigin-RevId: 618195554
 2024-03-22 09:05:39 -07:00
Kanglan Tang
3f1330805c Add an experimental build-only continuous cross compile build for MacOS x86 
PiperOrigin-RevId: 617611779
 2024-03-20 13:43:45 -07:00
Kanglan Tang
f0afc1b43d Add an experimental build-only continuous cross compile build for Linux Aarch64 
PiperOrigin-RevId: 613624879
 2024-03-07 10:17:43 -08:00
David Dunleavy
6928465b87 Add --use_clang and --clang_path options to build.py 
PiperOrigin-RevId: 603837975
 2024-02-02 18:20:44 -08:00
Nitin Srinivasan
b58772cdb4 Enable Bazel remote cache in macOS continuous builds 
PiperOrigin-RevId: 597100776
 2024-01-09 18:30:31 -08:00
jax authors
3d608b1850 Update CUDA to 12.3 in JAX/TF/XLA CIs 
This is updating CUDA to version 12.3. Related libraries (notably cuDNN) are also getting updated.

PiperOrigin-RevId: 596515360
 2024-01-08 01:24:32 -08:00
Jieying Luo
f25a51ee4f Change bazel config name "rbe_cpu_linux_py39" to "rbe_cpu_linux_py3.9"to be consistent with cuda bazel build configs. 
Change other python versions bazel config similarly.

PiperOrigin-RevId: 588100957
 2023-12-05 09:42:14 -08:00
Peter Hawkins
6fa506cb4a Build CUDA kernels for sm_50 instead of sm_52. 
This improves compatibility with older Maxwell cards, and it probably doesn't matter a whole lot for performance.

See: https://github.com/google/jax/issues/5723#issuecomment-1823093159
PiperOrigin-RevId: 585967281
 2023-11-28 07:16:01 -08:00
jax authors
dce6ab4548 Reverts 2aaa7559f96e4bb7b0271665bf386bf3ba22c451 
PiperOrigin-RevId: 584033001
 2023-11-20 08:23:40 -08:00
jax authors
2aaa7559f9 Reverts a80d35f4eda99bac42a1573ff5029e7035ebe71b 
PiperOrigin-RevId: 581285053
 2023-11-10 09:32:29 -08:00
jax authors
a80d35f4ed Replace gcc with clang compiler in CI nightly jobs. 
PiperOrigin-RevId: 581086810
 2023-11-09 18:30:03 -08:00
Jieying Luo
0ce7c7b7bd Register plugin profiler for TPU and remove --config=tpu/--enable_tpu in jaxlib. 
PiperOrigin-RevId: 580561059
 2023-11-08 09:40:28 -08:00
jax authors
59192b0dd1 Replace gcc with clang compiler in presubmit and postsubmit CI Kokoro jobs. 
PiperOrigin-RevId: 579032820
 2023-11-02 18:01:36 -07:00
Peter Hawkins
8d49f9a159 Reverts 9b1a656c1ef1f93f9e93eccb662de9bebe66b51a 
PiperOrigin-RevId: 576128882
 2023-10-24 07:08:34 -07:00
jax authors
9b1a656c1e Reverts 84c516974ab2b37169938d2d48a6c29a63c62c21 
PiperOrigin-RevId: 575891656
 2023-10-23 12:13:56 -07:00
Jieying Luo
0290150c4c Build jaxlib without PJRT GPU deps when plugin will be built. 
PiperOrigin-RevId: 573844805
 2023-10-16 09:59:07 -07:00
Yash Katariya
f5a1439168 Delete cuda 12.0.1 rbe configs since JAX doesn't support it anymore 
PiperOrigin-RevId: 573059967
 2023-10-12 17:52:11 -07:00
Peter Hawkins
d856ecc6fb Set RPATH, not RUNPATH in JAX CUDA builds. 
Fixes https://github.com/google/jax/issues/17497
 2023-10-12 09:38:10 -07:00
Peter Hawkins
efc18e4147 [JAX] Obtain NCCL via a stub, rather than linking it statically or dynamically. 
This shrinks the CUDA jaxlib wheel size by around 80MB.

PiperOrigin-RevId: 570554454
 2023-10-03 18:33:58 -07:00
Yash Katariya
1c90a5dfd0 Add cuda 12.2 support to JAX 
PiperOrigin-RevId: 568244230
 2023-09-25 09:40:12 -07:00
Yash Katariya
9ca133b5da Add python 3.12 support to JAX 
PiperOrigin-RevId: 567307333
 2023-09-21 07:56:05 -07:00
Peter Hawkins
de5aee8676 Fix Bazel build failures from long command lines on Windows. 
Fixes https://github.com/google/jax/issues/14950
 2023-09-07 17:36:04 -04:00
Peter Hawkins
64d0a36eab Don't enable the XLA Python --enable_tpu flag when building with --enable_tpu. 
This flag no longer does anything in XLA.

PiperOrigin-RevId: 559237161
 2023-08-22 15:12:46 -07:00
Peter Hawkins
1d4b10b775 Remove --distinct_host_configuration from Bazel flags. 
This flag does nothing under Bazel 6 and will be removed in Bazel 7.
 2023-07-11 11:38:05 -04:00
Peter Hawkins
bfa113ba60 Remove references to Python 3.8. 
Remove the old build scripts/Dockerfile, since they are unused and broken.

PiperOrigin-RevId: 542870354
 2023-06-23 08:48:57 -07:00
Peter Hawkins
119661ce6b Remove older plugin device integration. 
Users of this mechanism should migrate to the newer PJRT plugin registration mechanism (see the comments on discover_plugins() in this file).
 2023-06-14 15:26:58 -04:00
Peter Hawkins
cb33fdf3f7 Include SASS/PTX for Hopper GPUs. 2023-06-05 09:42:12 -04:00
Peter Hawkins
84c516974a Revert: Switch to using Clang as the default compiler. 
It appears this is causing deadlocks in multi-gpu tests.

PiperOrigin-RevId: 527706573
 2023-04-27 15:52:28 -07:00
Parker Schuh
782d90dc85 Switch to using Clang as the default compiler. 
PiperOrigin-RevId: 526815933
 2023-04-24 19:01:49 -07:00
Yash Katariya
86c0b36bfd Remove Cuda 11.4 support. JAX from 0.4.8 release will support cuda 11.8 and cuda 12 
PiperOrigin-RevId: 520056811
 2023-03-28 09:54:36 -07:00
Peter Hawkins
b62f114524 Add support for using pip-installed CUDA wheels. 
Add a currently undocumented jax[cuda11_pip] and jax[cuda12_pip] that depend on the pip CUDA wheels.
Add a currently undocumented jax[cuda11_local] and jax[cuda12_local] that avoid the CUDA wheel dependency.
 2023-03-26 12:35:00 +00:00
Yash Katariya
06d610dcd4 Add Cuda 12 build configs to bazelrc 
PiperOrigin-RevId: 516370916
 2023-03-13 17:46:22 -07:00
Peter Hawkins
172a831219 Switch JAX to use the OpenXLA repository. 2023-03-13 18:38:26 +00:00
Stella Laurenzo
c1e13bdf3f A few developer workflow enhancements for working with jaxlib. 
It seems to me that jaxlib development must be mostly happening on CI, because some basics are pretty essential. Here are a few things I've been typing/carrying for a while in my flow:

* Add .bazelrc.user to .gitignore so it doesn't accidentally get checked in.
* Add configs for 'debug_symbols' and 'debug' that make some things minimally workable under a debugger (or to get backtraces, etc).
* Add `--force-reinstall` to the copy/paste command to update a built jaxlib wheel (without this, if you are iterating, it fairly quietly does nothing).
 2023-02-10 21:03:21 -08:00
Peter Hawkins
ec56d71d01 Drop support for NVIDIA Kepler series GPUs in jaxlib builds. 2023-02-10 14:15:15 -05:00
Yash Katariya
71569e1d16 Remove the specialized sm versions for testing. It caused release wheels to ship without correct sm versions compiled in. 
PiperOrigin-RevId: 494854373
 2022-12-12 16:02:46 -08:00
Yash Katariya
ea63baf1ad Remove python 3.7 from testing and release builds 
PiperOrigin-RevId: 491752393
 2022-11-29 15:26:46 -08:00
Yash Katariya
03aa266e45 Build GPU wheel only for Tesla (t4) when running on RBE. THis should in theory speed up the builds 
PiperOrigin-RevId: 490553589
 2022-11-23 11:57:24 -08:00
Yash Katariya
51e4b017bb Use the cuda 11.8 image for CPU builds to reduce the docker image churn from the JAX side in TF's RBE cluster 
PiperOrigin-RevId: 490522647
 2022-11-23 09:33:17 -08:00
Yash Katariya
6d11567142 Add cuda 11.8 configs to .bazelrc 
PiperOrigin-RevId: 489591457
 2022-11-18 16:55:28 -08:00
Yash Katariya
d77cccff43 Update docker images to add python 3.11 support. 
PiperOrigin-RevId: 485488215
 2022-11-01 20:35:38 -07:00
Rohit Santhanam
663441007c [ROCm] Added gfx90a and gfx1030. 2022-10-27 16:57:06 +00:00
Peter Hawkins
aa6445e6f5 Set --tsl_link_protobuf=true in .bazelrc to unbreak jaxlib build. 2022-10-13 13:50:09 +00:00
jax authors
254dc24a8b Merge pull request #11961 from jakeh-gc:plugin_device 
PiperOrigin-RevId: 476363760
 2022-09-23 07:29:17 -07:00
jax authors
fd90f40c45 Merge pull request #12443 from cloudhan:fix-mlir-chlo-stablehlo-symbols 
PiperOrigin-RevId: 475808753
 2022-09-21 06:12:44 -07:00
Cloud Han
dc414c8635 Switch from experimental_strict_action_env to incompatible_strict_action_env to avoid deprecation warning 2022-09-21 17:26:21 +08:00