<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.

## 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`. At the moment you will have to use a Git dependency while we work on our first release:

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

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

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

### 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.

## Configuration

### In projects

EXLA works both as a backend for Nx tensors and an optimized `Nx.Defn` compiler. To enable both globally, add a `config/config.exs` (or `config/ENV.exs`) with the following:

```elixir
import Config
config :nx, :default_backend, EXLA.Backend
config :nx, :default_defn_options, [compiler: EXLA]
```

You can also use cuda/rocm/tpu as the target by setting `:client` option in both configuration:

```elixir
config :nx, :default_backend, {EXLA.Backend, client: :cuda}
config :nx, :default_defn_options, [compiler: EXLA, client: :cuda]
```

To use GPUs/TPUs, you must also set the appropriate value for the [`XLA_TARGET`](https://github.com/elixir-nx/xla#xla_target) environment variable. For CUDA, setting `ELIXIR_ERL_OPTIONS="+sssdio 128"` is also required on more complex operations to increase CUDA's compiler stack size.

### In scripts/notebooks

The simplest way to configure EXLA in notebooks is by calling:

```elixir
EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])
```

The above will set EXLA both as backend and as a compiler. To use GPUs/TPUs, you must also set the appropriate value for the [`XLA_TARGET`](https://github.com/elixir-nx/xla#xla_target) environment variable.

## Contributing

### Building locally

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

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

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.