86 Commits

Author SHA1 Message Date
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
Cloud Han
a6acce58e0 Build on Windows
1. Build on Windows

2. Fix OverflowError

    When calling `key = random.PRNGKey(0)` OverflowError: Python int too
    large to convert to C long for casting value 4294967295 (0xFFFFFFFF)
    from python int to int32.

3. fix file path in regex of errors_test

4. handle ValueError of os.path.commonpath
2020-11-19 23:33:06 +08:00
Peter Hawkins
b222e5e05f Update bazel version to 3.1.0. 2020-11-17 15:50:30 -05:00
Peter Hawkins
2c6f932e0e Add Python 3.9 support to jaxlib build. 2020-11-09 16:25:58 -05:00
Srijan Saurav
40e20242db
Fix code quality issues (#4302)
Changes:
- Fix unnecessary generator
- Iterate dictionary directly instead of calling .keys()
- Remove global statement at the module level
- Use list() instead of a list comprehension
- Use with statement to open the file
- Merge isinstance calls
2020-09-17 09:21:18 -07:00
Peter Hawkins
a141cc6e8d
Make CUDA wheels manylinux2010 compliant, add CUDA 11, drop CUDA 9.2 (#3555)
* Use dynamic loading to locate CUDA libraries in jaxlib.

This should allow jaxlib CUDA wheels to be manylinux2010 compliant.

* Tag CUDA jaxlib wheels as manylinux2010.

Drop support for CUDA 9.2, add support for CUDA 11.0.

* Reorder CUDA imports.
2020-06-25 14:37:14 -04:00