<h1><img src="https://github.com/elixir-nx/nx/raw/main/exla/exla.png" alt="EXLA" width="350"></h1>

[Google's XLA](https://www.tensorflow.org/xla/) (Accelerated Linear Algebra) compiler/backend for Nx. It supports just-in-time (JIT) compilation to GPU (both CUDA and ROCm) and TPUs.

[See the documentation](https://hexdocs.pm/exla).

## Installation

In order to use `EXLA`, you will need Elixir installed. Then create an Elixir project via the `mix` build tool:

```
$ mix new my_app
```

Then you can add `EXLA` as dependency in your `mix.exs`:

```elixir
def deps do
  [
    {:exla, "~> 0.5"}
  ]
end
```

If you are using Livebook or IEx, you can instead run:

```elixir
Mix.install([
  {:exla, "~> 0.5"}
])
```

Once installed, you must configure Nx to use EXLA by default. [Check out the "Configuration" section in the docs](https://hexdocs.pm/exla/EXLA.html#module-configuration) to learn how to do so.

### XLA binaries

EXLA relies on the [XLA](https://github.com/elixir-nx/xla) package to provide the necessary XLA binaries. Whenever possible it tries to download precompiled builds, but you may need to build from source if there is no version matching your target environment. For more details, including GPU/TPU support see [the usage section](https://github.com/elixir-nx/xla#usage).

### Common installation issues

  * Missing Dependencies
    * Some Erlang installs do not include some of the dependencies needed to compile the EXLA NIF. You may need to install `erlang-dev` separately.
  * Incompatible protocol buffer versions
    * Error message: "this file was generated by an older version of protoc which is incompatible with your Protocol Buffer headers".
    * If you have `protoc` installed on your machine, it may conflict with the `protoc` precompiled inside XLA. Uninstall, unlink, or remove `protoc` from your path to continue.

### Usage with Nerves

For cross-compilation, you need to [set your `XLA_TARGET_PLATFORM` variable](https://github.com/elixir-nx/xla#xla_target_platform) to the correct target platform value (i.e. `aarch64-linux-gnu` for the Raspberry Pi 4).

## Troubleshooting

EXLA uses NIFs (C-interface code called from Elixir) for part of its functionality.
If for any reason these fail to compile or load, troubleshooting can be tricky.

We recommend following the steps below:

  1. If the error appeared after upgrading EXLA, ensure that you have the proper dependency versions given by [XLA](https://github.com/elixir-nx/xla). Afterwards, compile with `mix compile` after setting `EXLA_FORCE_REBUILD` to clean up cached files:
     * `EXLA_FORCE_REBUILD=partial`: Removes the only the libexla.so caches (both local and global ones).
     * `EXLA_FORCE_REBUILD=true`: Removes the libexla.so caches but also removes the intermediate `.o` compilation artifacts retained from previous builds.

     Additional notes on compilation:
     * Besides the XLA dependency versions, ensuring `gcc` (or your compiler of choice), `libc` and `make` are compatible is also important.
     * Remember to save the compilation logs from this step for further debugging.
     * It is a good idea to save the `cache/<version>/libexla.so` file so that the team can inspect its contents if needed.
  2. If the error persists, look for the `** (RuntimeError) Failed to load NIF library.` exception on application start-up.
    This exception should provide more information on what's the issue when loading the NIF. Share these logs in an issue on GitHub
    so that the Nx team can investigate further.


## Contributing

### Building locally

EXLA is a regular Elixir project, therefore, to run it locally:

```shell
mix deps.get
mix test
```

By default, EXLA passes `["-jN"]` as a Make argument, where `N` is `System.schedulers_online() - 2`, capped at `1`. `config :exla, :make_args, ...` can be used to override this default setting.

In order to run tests on a specific device, use the `EXLA_TARGET` environment variable, which is a dev-only variable for this project (it has no effect when using EXLA as a dependency). For example, `EXLA_TARGET=cuda` or `EXLA_TARGET=rocm`. Make sure to also specify `XLA_TARGET` to fetch or compile a proper version of the XLA binary.

### Building with Docker

The easiest way to build is with [Docker](https://docs.docker.com/get-docker/). For GPU support, you'll also need to set up the [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-docker).

To build, clone this repo, select your preferred Dockerfile, and run:

```shell
docker build --rm -t exla:host . # Host Docker image
docker build --rm -t exla:cuda10.2 . # CUDA 10.2 Docker image
docker build --rm -t exla:rocm . # ROCm Docker image
```

Then to run (without Cuda):

```shell
docker run -it \
  -v $PWD:$PWD \
  -e TEST_TMPDIR=$PWD/tmp/bazel_cache \
  -e BUILD_CACHE=$PWD/tmp/xla_extension_cache \
  -w $PWD \
  --rm exla:host bash
```

With CUDA enabled:

*Note: XLA_TARGET should match your CUDA version. See: https://github.com/elixir-nx/xla#xla_target*

```shell
docker run -it \
  -v $PWD:$PWD \
  -e TEST_TMPDIR=$PWD/tmp/bazel_cache \
  -e BUILD_CACHE=$PWD/tmp/xla_extension_cache \
  -e XLA_TARGET=cuda102 \
  -e EXLA_TARGET=cuda \
  -w $PWD \
  --gpus=all \
  --rm exla:cuda10.2 bash
```

With ROCm enabled:

```shell
docker run -it \
  -v $PWD:$PWD \
  -e TEST_TMPDIR=$PWD/tmp/bazel_cache \
  -e BUILD_CACHE=$PWD/tmp/xla_extension_cache \
  -e XLA_TARGET=rocm \
  -e EXLA_TARGET=rocm \
  -w $PWD \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --rm exla:rocm bash
```

Inside the container you can interact with the API from IEx using:

```shell
iex -S mix
```

Or you can run an example:

```shell
mix run examples/regression.exs
```

To run tests:

```shell
mix test
```

## License

Copyright (c) 2020 Sean Moriarity

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at [http://www.apache.org/licenses/LICENSE-2.0](http://www.apache.org/licenses/LICENSE-2.0)

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.