Add a new build option --enable_mkl_dnn that enables MKLDNN contraction kernels in XLA. This leads to significant performance improvements for XLA's dot operator. Enable MKL-DNN by default.
Update XLA version to include MKL-DNN build fix.
Also add a new --enable_march_native build option that turns on -march=native. This is unlikely to have a significant performance impact since XLA JIT-compiles most of its code. Leaving this off by default because it also generates code unlikely to run across a wide selection of architectures and so is unsuitable for building pip wheels.
Without this change build script fails with the following message:
```
Encountered error while reading extension file 'tensorflow/workspace.bzl': no such package '@org_tensorflow//tensorflow': The native http_archive rule is deprecated. load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") for a drop-in replacement.
```
This PR simply implements the suggestion in the error. I was able to successfully build after applying the change.
This makes initial builds cheaper (since we don't need to build some files in separate host and target configurations) but may make switching between build configurations more expensive (since we can share less work). The build script should optimize for the former.
Significant changes:
* Mac OS X support.
* build script is in Python, not shell.
* build configuration is passed via flags, not environment variables.
* build script configures TF itself, and does not require explicitly checking out the TF git repository and running its configure script. Changes the TF dependency in the Bazel workspace to be an http_archive(), rather than a local checkout of TF.
* rather than trying to guess the path for Bazel-generated XLA artifacts, use a sh_binary() to perform installation of the built artifacts in to the JAX source tree. Bazel's runfiles mechanism is the supported route to find build artifacts.
* downloads Bazel in Python and checks its SHA256 before running it, rather than running an untrusted binary from the internet.
* intentionally does not delete the Bazel cache or Bazel after building.
Example of new build interaction:
Building without CUDA on Mac or Linux:
$ cd jax
$ python3 build.py (or python2 build.py if you want a Python 2 build)
_ _ __ __
| | / \ \ \/ /
_ | |/ _ \ \ /
| |_| / ___ \ / \
\___/_/ \_\/_/\_\
Starting local Bazel server and connecting to it...
Bazel binary path: /Users/xyz/bin/bazel
Python binary path: /Library/Frameworks/Python.framework/Versions/3.7/bin/python3
CUDA enabled: no
Building XLA and installing it in the JAX source tree...
...
Example of building with CUDA enabled on Linux:
$ python3 build.py --enable_cuda --cudnn_path=/usr/lib/x86_64-linux-gnu/
... as before, except ...
CUDA enabled: yes
CUDA toolkit path: /usr/local/cuda
CUDNN library path: /usr/lib/x86_64-linux-gnu/
...
PiperOrigin-RevId: 222868835
