mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
11,379 Commits 503 Branches 263 Tags
Commit Graph

91 Commits

Author SHA1 Message Date
Yash Katariya
6ba9fb699d Upgrade the bazel version to 5.1.1 
PiperOrigin-RevId: 441338363
 2022-04-12 17:48:09 -07:00
Yash Katariya
aa5d6b4a58 Fix the breakage by including --experimental_cc_shared_library as done by TF. 
PiperOrigin-RevId: 438746867
 2022-03-31 23:07:42 -07:00
Reza Rahimi
a0d9d81f92 Update JAX to use new math libraries in ROCm-5.0. 2022-03-01 20:02:15 +00:00
Peter Hawkins
2e0cfe8e42 Update the list of default CUDA capabilities used for wheel builds to match build.py. 2022-02-15 09:23:28 -05:00
Peter Hawkins
2388e353da Increase bazel version to 5.0.0 to match TensorFlow 
(8871926b0a).
 2022-01-28 21:11:02 +00:00
Peter Hawkins
04369a3588 Drop support for NumPy 1.18. 
Per NEP-29, we can drop NumPy 1.18 support on Dec 22, 2021.

The next NumPy deprecation will be 1.19 on Jun 21, 2022.

PiperOrigin-RevId: 419651428
 2022-01-04 12:11:38 -08:00
Peter Hawkins
66823d1392 Include compute capability 8.0 SASS in jaxlib wheels. 
Drop compute capability 6.1 to avoid growing the wheel size.

Also fix an unrelated build error due to a gcc warning in boringssl.
 2021-12-14 14:27:19 -05:00
Peter Hawkins
ffb7ec1651 Update Bazel to 4.2.1. 
Fixes #8573
 2021-12-02 09:11:38 -05:00
Peter Hawkins
9212d5c83b Print the bazel version from build.py. 
Increment the minimum version that build.py checks for to 3.7.2.
 2021-11-03 12:12:54 -04:00
Peter Hawkins
256e7220ff [JAX] Fix pylint errors. 
* trailing-whitespace
* dangerous-default-value. None of these appear to be bugs in practice, but the potential for accidentally mutating the default value is there, and the cost of avoiding the problem is small.
* invalid-envvar-default. Pass strings as getenv() defaults.
* unnecessary-semicolon. Use tuples instead for this one-liner.
* invalid-hash-returned. Raise an exception rather than asserting false.
* pointless-string-statement. Use comments instead.
* unreachable. Use @unittest.skip() decorator rather than raising as first line in test.
* logging-not-lazy. Make the logging lazy.
* bad-format-string-type. Use f-string instead.
* subprocess-run-check. Pass check=...

PiperOrigin-RevId: 400858477
 2021-10-04 17:54:46 -07:00
brett koonce
9c5009efd5 tweak cuda/rocm targets 
Closes #7955.
 2021-09-18 14:29:13 -05:00
yashkatariya
d0acd9f343 Add flags to configure the cuda_compute_capability and rocm_amd_targets 2021-09-17 08:43:25 -07:00
Peter Hawkins
94f97b920f Refactor JAX CPU kernels to make them usable from C++. 
Most of the work here is porting the LAPACK interface from Cython to plain C++. This is something I wanted to do anyway to make use of C++ templating facilities: the code is noticeably shorter in C++.

This change removes the only use of Cython in JAX. It also removes the need for a build-time dependency on Scipy, which we only needed for Cython cimport reasons.

When using C++, we most likely do not want to fetch LAPACK and BLAS kernels from Python. Therefore we add another option: we define the LAPACK functions we need using weak symbols where supported; the user can then simply link against LAPACK to provide the necessary symbols.

Added a jaxlib:cpu_kernels module to facilitate using the JAX CPU kernels from C++.

PiperOrigin-RevId: 394705605
 2021-09-03 10:03:54 -07:00
Yash Katariya
7dca05066a Remove build --strategy=Genrule=standalone since this makes genrule run locally instead of on remote RBE clusters leading to errors that are very hard to debug. 
Add this back in build.py because that script is used for building jaxlib locally.

PiperOrigin-RevId: 391603481
 2021-08-18 13:56:43 -07:00
Reza Rahimi
f454f6b7b8 fix rocm_amdgpu_targets for rocm 2021-08-17 22:13:29 +00:00
Yash Katariya
2d1854a8ba Create a .bazelrc file that is the base for all the builds. The current build.py workflow will not be affected since this .bazelrc file will be overridden. I am going to change that workflow in the coming CLs. 
PiperOrigin-RevId: 390003558
 2021-08-10 16:17:58 -07:00
jax authors
1a2ddc055d Merge pull request #7419 from cloudhan:win-builder-prepare 
PiperOrigin-RevId: 389938513
 2021-08-10 11:42:48 -07:00
yashkatariya
b24574f743 Fix the regex 2021-08-05 13:25:06 -07:00
yashkatariya
86aaf80dce Use bazel --version 2021-08-05 09:45:39 -07:00
Cloud Han
be1705306c fix bazel auto download 2021-07-31 19:13:16 +08:00
Cloud Han
d4349a42b5 workaround weird issue due to non-standard cuda path 2021-07-30 19:25:33 +08:00
Cloud Han
86e37d996a some fs support now can be disabled without compile error 2021-07-30 19:23:18 +08:00
Peter Hawkins
6e9169d100 Drop support for NumPy 1.17. 2021-07-29 09:18:01 -04:00
Peter Hawkins
6c08702489 Add support for ppc64le cross-compilation on Linux. 
Use Bazel 4.1.0 unconditionally on all platforms.
 2021-07-23 10:39:02 -04:00
Cloud Han
2d321c26e6 Use TF_CUDA_PATHS 
CUDA_TOOLKIT_PATH and CUDNN_INSTALL_PATH are deprecated, see TF 2.0
release notes for more information
 2021-07-18 22:55:34 +08:00
jax authors
6aa20d8f8f Merge pull request #7294 from hawkinsp:py36 
PiperOrigin-RevId: 384994957
 2021-07-15 13:19:23 -07:00
Peter Hawkins
94446ff757 Drop Python 3.6 support. 
Per the deprecation policy (https://jax.readthedocs.io/en/latest/deprecation.html),
Python 3.6 support has been due for removal since June 23, 2020.
 2021-07-15 14:20:29 -04:00
Peter Hawkins
f5c61a892a Add support for cross-compiling jaxlib for Mac ARM. 2021-07-15 10:37:53 -04:00
Peter Hawkins
7d2aec105f Add an option to disable NCCL. 2021-07-13 09:10:29 -04:00
Qiao Zhang
18c7610e96 Catch CalledProcessError in build script. 2021-06-22 13:36:14 -07:00
Peter Hawkins
b130257ee1 Drop support for NumPy 1.16. 2021-06-11 09:03:09 -04:00
akbir khan
dc81610bb6 updated to official bazel.4.1.0 2021-05-24 21:40:08 +01:00
Peter Hawkins
dacf31f202 Check for NumPy and SciPy versions during jaxlib builds. 2021-05-24 12:39:37 -04:00
Peter Hawkins
86887a21e2 Add Mac M1 support to build.py. 2021-05-18 21:53:31 -04:00
Peter Hawkins
7d4551dd45 Add support for arm64 builds to build.py and build_wheel scripts. 2021-05-10 11:08:09 -04:00
Skye Wanderman-Milne
7b7a2a1ba8 Don't build with Cloud TPU support for GPU or Mac wheels. 
This should make our builds simpler and less failure-prone.
 2021-03-08 16:19:19 -08:00
Christian Sigg
c0e42e59eb
Switch to new flag to enable CUDA support in TF 2021-03-03 08:55:08 +01:00
Situ Yi
77905ef5fd
Update build.py with Cloud TPU support 
Update build.py to have Cloud TPU support by default.
b/180656002
 2021-02-19 12:57:57 -08:00
jax authors
15359e018d Merge pull request #5738 from hawkinsp:avx 
PiperOrigin-RevId: 357785312
 2021-02-16 12:53:41 -08:00
Skye Wanderman-Milne
3caf52d54a Fix build.py version comparison 2021-02-16 11:32:50 -08:00
Peter Hawkins
a27545497b Enable AVX by default in release builds. 
TensorFlow also enables AVX by default, so this seems like a reasonable baseline CPU feature to expect.
 2021-02-16 14:12:01 -05:00
jax authors
5fd3deb0be Merge pull request #5580 from daskol:search-bazel-in-path 
PiperOrigin-RevId: 357755577
 2021-02-16 10:49:33 -08:00
Skye Wanderman-Milne
dd513c1585
Small simplification 2021-02-16 08:39:32 -08:00
Daniel Bershatsky
96193cef12 Simplify bazel version check 2021-02-13 16:32:21 +03:00
Skye Wanderman-Milne
74db363830 Bump required bazel version to 3.7.2 
Untested, but I'm about to build a new jaxlib and will make any fixes then.
 2021-02-11 09:07:04 -08:00
Daniel Bershatsky
c71a4f5871 Try to download bazel as the last resort 
Resolution order of paths to bazel binary is as follows.

1. Use --bazel_path command line option.
2. Search bazel binary in PATH environment variable.
3. Download required bazel release.
 2021-02-01 15:15:38 +03:00
Clemens Giuliani
330e69dd30 Add support for ROCm 2020-12-06 22:55:53 +01:00
Cloud Han
ea340eed93 use strict action env on windows to avoid constant full rebuilding 2020-11-24 23:44:26 +08:00
Peter Hawkins
c06ead6b04 Change jaxlib build rules to build a wheel, rather than writing output to the source directory. 2020-11-20 11:47:00 -05:00
Peter Hawkins
1fcbd2f083 Add -Wno-stringop-truncation to build flags on Linux. 
Works around https://github.com/tensorflow/tensorflow/issues/39467 for gcc 10+ builds.
 2020-11-19 22:50:08 -05:00