EXLAでcudaが動作する環境の構築と実際にNxで動かしたデモ

※一年前に作成て下書きに眠っていたので公開しておきます

# 環境構築

## GPUを購入する

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/906ddf81-b303-a549-a1f6-ac86b5507ff3.png)

## PCに取り付ける

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/1262a05e-e4f3-4d69-497a-87512bd38ff3.png)

## Windows用Driverインストール

[Run Linux GUI apps on the Windows Subsystem for Linux](https://learn.microsoft.com/en-us/windows/wsl/tutorials/gui-apps) に従ってvGPUのドライバ（Windows側）をインストールする。

NVidiaの場合は、以下からダウンロード
https://developer.nvidia.com/cuda/wsl

https://learn.microsoft.com/en-us/windows/wsl/tutorials/gui-apps
※[ドキュメント](https://learn.microsoft.com/en-us/windows/wsl/tutorials/gui-apps)にvGPUのドライバーをインストールする必要があると書いてあるので必要なのだと思いますが、本当に必要なのかは試していません。
もしかしたら、WindowsのInboxのドライバだけで動くかも？このドライバをインストールしなくても動くかどうか、試した方いらっしゃったらコメントいただけるありがたいです。

## WSL2の環境構築(Not docker)

```bash
sudo apt-key del 7fa2af80
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/3bf863cc.pub
sudo add-apt-repository 'deb https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/ /'
sudo apt-get update
sudo apt-get -y install cuda=11.1.0-1 #バージョン指定する

```

## WSL2の環境構築(docker)

https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#setting-up-nvidia-container-toolkit


## デバイス認識確認
nvidia-smiコマンドを実行して、デバイス状態を表示できるか確認

```
$ nvidia-smi
Sun Oct  9 20:32:05 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.75       Driver Version: 517.40       CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0 Off |                  N/A |
|  0%   31C    P8    12W / 170W |  11800MiB / 12288MiB |      4%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      7205      C   /beam.smp                       N/A      |
+-----------------------------------------------------------------------------+
```

# アプリケーション編

## Elixir環境の動作確認

iexから次のコマンドを実行

```
Mix.install([
  {:nx, "~> 0.3.0"},
  {:exla, "~> 0.3.0"},
],
config: [
          nx: [
            default_backend: EXLA.Backend,
            default_defn_options: [compiler: EXLA],
          ]
        ],
system_env: [
    XLA_TARGET: "cuda111"
  ],
  # XLA_TARGETの変更で、exlaを再コンパイルしたい時はtrueにする
  force: false,
)
Nx.add(Nx.tensor([1]), Nx.tensor([1]))
```

## mixの設定箇所

環境変数XLA_TARGETを設定。.bashrcに記載するなど。

```bash:.bashrc
export XLA_TARGET=cuda111
```

depsに、exlaを追加

```elixir:mix.exs
  defp deps do
    [
      {:axon, "~> 0.2.0"}, # axonを使う場合
      {:exla, "~> 0.3.0"},
      {:nx, "~> 0.3.0"},
    ]
  end
```
configに追加
```elixir:config/config.exs
import Config
config :nx, :default_backend, EXLA.Backend
config :nx, :default_defn_options, compiler: EXLA
```

## プログラムの対応

GPUで処理する内容をdefnに記述する

```elixir:exla101.ex
defmodule Exla101 do
  import Nx.Defn
  defn add do
    a = Nx.tensor([[1,2,3],[4,5,6]])
    b = Nx.tensor([[1,1,1],[2,2,2]])
    Nx.add(a,b)
  end

```