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
18,626 Commits 508 Branches 263 Tags
Commit Graph

68 Commits

Author SHA1 Message Date
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
Yash Katariya
a2e05b0df5 Try Cuda 11.2 build in presubmit for GPU. Presubmit uses 11.2 and CI uses 11.1 
PiperOrigin-RevId: 473071603
 2022-09-08 13:34:03 -07:00
Jake
21f82c6c0d Use the pjrt plugin device client. 2022-08-17 14:34:07 +01:00
Peter Hawkins
184f8ab723 Import .bazelrc.user to allow users to override bazel options easily. 2022-08-14 15:46:12 -04:00
Rohit Santhanam
080cf47002 [ROCm] Fixes for compilation failures caused by compiler changes in ROCm Tensorflow fork. 2022-06-29 14:34:08 +00:00
Peter Hawkins
69bda69fb6 Bump minimum Mac OS version to 10.14 (Mojave). 
It turns out that the support for C++17 is partial in 10.12, and in particular absl::optional and std::optional are not the same thing under 10.12. Increment to 10.14 which is the lowest version that builds successfully with absl::optional == std::optional.

See: 89cdaed655/absl/base/config.h (L528)
Strictly speaking, we could allow 10.13, but not without updating ABSL in the TF repository to incorporate c86347d4ce which fixes the version detection test to permit 10.13 as well.
 2022-06-01 20:32:22 -04:00
Nicholas Junge
7f7358cb58 Fix macOS builds by updating the minimum OSX version target 
This commit bumps the minimum macOS version target to 10.12, from 10.9
previously. The reason is that a newer LLVM version, which the XLA
compiler depended on, used `std::shared_mutex`, which is only available
starting at macOS 10.12.

The fix here consists of setting the minimum version target flag that
bazel consumes for the build to 10.12.
 2022-05-31 15:14:48 +02:00
Yash Katariya
dd261bde59 Switch to using C++17 
PiperOrigin-RevId: 450559689
 2022-05-23 17:15:00 -07:00
Peter Hawkins
ea868b4ffe Update .bazelrc to reflect GPU/TPU build changes. 
PiperOrigin-RevId: 449306362
 2022-05-17 13:48:55 -07:00
Peter Hawkins
9fb9e12169 Don't include PTX for older GPU generations. 
See: https://github.com/tensorflow/tensorflow/pull/55613

For a CUDA build at head with the default compute capabilities, reduces wheel size from 141MB to 112MB.

Don't redundantly specify default compute capabilities in .bazelrc and in
build.py.
 2022-05-02 20:27:37 +00:00
Yash Katariya
5fdad0ebf5 Roll forward manylinux2014 builds after fixes. 
PiperOrigin-RevId: 440589273
 2022-04-09 10:19:44 -07:00