【2022/10月版】WSL2のElixir、NxでGPUを使ってみた

# はじめに

WSL2で、Nxを使ってますが、ここでGPUを使いたい。

この環境構築時の手順の記録です。

### 環境
OS Windows 11
WSL2で動作しているOS　Ubuntu 20.04 LTS
GPU　RTX3060


[enabling-gpu-acceleration-on-ubuntu-on-wsl2-with-the-nvidia-cuda-platform](https://ubuntu.com/tutorials/enabling-gpu-acceleration-on-ubuntu-on-wsl2-with-the-nvidia-cuda-platform#1-overview)に従ってインストールしました。
nVidiaから、WSL2用のパッケージが公開されているので、特に問題なく導入できました。

XLAがcuda11.1でないと[動作しなかった](https://qiita.com/GeekMasahiro/items/1a9226101df99948db9c)ので、WSL側でcudaパッケージをインストールする時に、バージョン指定をしました。


# Windows 側のDriverのインストール

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

[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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/10292216-d977-e03a-59c1-bf3b73c30a2c.png)

Product Typeなどを選び、DEARCHをクリック、
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/af2f6e1d-da9c-c9ac-c771-b8e3f08cf641.png)

ダウンロードをクリック

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/4f85a5da-a397-7e30-735b-d622c5ed79bb.png)

ダウンロードしたファイルを実行

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/94d0de30-1e8c-af92-c2ef-9dff53c8f353.png)

NVIDIAグラフィックスドライバー　だけでよいらしいのでこちらを選択
同意して継続

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/8dfba0f7-19a9-f896-a364-94ca217fb76d.png)

高速を選択

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/91270704-9f9c-f0c5-fbbd-78e3d9b4eea3.png)

PCを再起動

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/58f06ca2-7efb-437b-eba9-2fdec96f8c62.png)



# WSL2(Ubuntu側)の設定

## cudaパッケージインストール

[ubuntuのドキュメント](https://ubuntu.com/tutorials/enabling-gpu-acceleration-on-ubuntu-on-wsl2-with-the-nvidia-cuda-platform#3-install-nvidia-cuda-on-ubuntu)に従ってインストールします。
wsl-ubuntuのcudaパッケージを導入するところがポイントです。

XLAのコンパイル済みのXLAがcuda 11.1を想定しているので、バージョンを11.1.0-1指定してインストールします。
バージョンを指定しない場合、最新のパッケージがインストールされます。Nxを実行時にエラーがでて動作しませんでした。


```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 #バージョン指定する

```

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      |
+-----------------------------------------------------------------------------+
```

nvidia-smi.exeの方を実行するとWindows側で動作しているプロセスを確認する事ができます。

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

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      4008    C+G   ...ge\Application\msedge.exe    N/A      |
|    0   N/A  N/A      4488    C+G   ...icrosoft VS Code\Code.exe    N/A      |
|    0   N/A  N/A      7292    C+G   ...2txyewy\TextInputHost.exe    N/A      |
|    0   N/A  N/A      9592    C+G   ...ontend\Docker Desktop.exe    N/A      |
|    0   N/A  N/A     10824    C+G   ...\PowerToys.FancyZones.exe    N/A      |
|    0   N/A  N/A     11480    C+G   C:\Windows\explorer.exe         N/A      |
|    0   N/A  N/A     12008    C+G   ...lPanel\SystemSettings.exe    N/A      |
|    0   N/A  N/A     12508    C+G   ...artMenuExperienceHost.exe    N/A      |
|    0   N/A  N/A     12532    C+G   ...n1h2txyewy\SearchHost.exe    N/A      |
|    0   N/A  N/A     13764    C+G   ...y\ShellExperienceHost.exe    N/A      |
|    0   N/A  N/A     14716    C+G   ...ext\CNext\AMDRSSrcExt.exe    N/A      |
|    0   N/A  N/A     16012    C+G   ...e\PhoneExperienceHost.exe    N/A      |
|    0   N/A  N/A     16696    C+G   ...ncher\CiscoCollabHost.exe    N/A      |
|    0   N/A  N/A     17536    C+G   ...ser\Application\brave.exe    N/A      |
|    0   N/A  N/A     18492    C+G   ...v10z8vjag6ke6\HP.MyHP.exe    N/A      |
|    0   N/A  N/A     20980    C+G   ...9cpj4tp0\Amazon Music.exe    N/A      |
|    0   N/A  N/A     21048    C+G   ...he8kybcnzg4\app\Slack.exe    N/A      |
|    0   N/A  N/A     21500    C+G   ...\CNext\RadeonSoftware.exe    N/A      |
|    0   N/A  N/A     21752    C+G   ...qxf38zg5c\Skype\Skype.exe    N/A      |
|    0   N/A  N/A     22976    C+G   ...370.34\msedgewebview2.exe    N/A      |
|    0   N/A  N/A     28336    C+G   ...8bbwe\WindowsTerminal.exe    N/A      |
|    0   N/A  N/A     28368    C+G   ...cw5n1h2txyewy\LockApp.exe    N/A      |
|    0   N/A  N/A     29648    C+G   ...wekyb3d8bbwe\Video.UI.exe    N/A      |
|    0   N/A  N/A     33524    C+G   ...370.37\msedgewebview2.exe    N/A      |
+-----------------------------------------------------------------------------+
```

## サンプルアプリケーションで動作を確認（オプション）

サンプルプログラムをコンパイルして実行
サンプルプログラムが大量にあるので、コンパイルに時間がかかります。必須ではありません。

```bash
git clone https://github.com/nvidia/cuda-samples
cd cuda-samples
make
```
libglutが無いというエラーが発生したので、ライブラリーを追加して、再度makeを実行

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/04f54d26-94d7-ced1-ba97-b2c3d43723e1.png)


```bash
sudo apt-get -y install freeglut3-dev
make

```

実効コマンド

```
$ ./bin/x86_64/linux/release/deviceQuery
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/532235/e98bdf8e-3ce6-5e41-7e40-b3fef8c6097a.png)

# EXLAを有効化

mix.exsにexlaを追加
```elixir:mix.exs
  defp deps do
    [
      {:axon, "~> 0.2.0"},
      {:exla, "~> 0.3.0"},
      {:nx, "~> 0.3.0"},
      {:kino_vega_lite, "~> 0.1.3"},
      {:jason, "~> 1.2"}
    ]
  end
```

次の設定をconfig.exsに追加
```Elixir:config/config.exs
import Config
config :nx, :default_backend, EXLA.Backend
config :nx, :default_defn_options, compiler: EXLA
```

環境変数を設定して、iexを起動。depsをcleanしないとうまく行かない場合があったので念のためそのコマンドも書いておきます。
```
export XLA_TARGET=cuda111
export XLA_BUILD=false
mix deps.clean --all
mix deps.get
iex -S mix
```

うまく行っていれば、XLAドライバー`xla_extension-x86_64-linux-cuda111.tar.gz`がダウンロードされて使用されます。

実行結果

```
$ iex -S mix
Erlang/OTP 25 [erts-13.0.4] [source] [64-bit] [smp:16:16] [ds:16:16:10] [async-threads:1] [jit:ns]

～～～省略～～～
==> elixir_make
Compiling 1 file (.ex)
Generated elixir_make app
==> xla
Compiling 2 files (.ex)
Generated xla app

20:14:44.654 [info] Found a matching archive (xla_extension-x86_64-linux-cuda111.tar.gz), going to download it
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                                                                                                Dload  Upload   Total   Spent    Left  Speed
  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0
100  204M  100  204M    0     0  5500k      0  0:00:37  0:00:37 --:--:-- 3780k

20:15:22.641 [info] Successfully downloaded the XLA archive
==> exla
Unpacking /home/masa/.cache/xla/0.3.0/cache/download/xla_extension-x86_64-linux-cuda111.tar.gz into /home/masa/deeplearning4/4.2/gridworld/deps/exla/cache
g++ -fPIC -I/home/masa/.asdf/installs/erlang/25.0.4/erts-13.0.4/include -Icache/xla_extension/include -O3 -Wall -Wno-sign-compare -Wno-unused-parameter -Wno-missing-field-initializers -Wno-comment -shared -std=c++14 c_src/exla/exla.cc c_src/exla/exla_nif_util.cc c_src/exla/exla_client.cc -o cache/libexla.so -Lcache/xla_extension/lib -lxla_extension -Wl,-rpath,'$ORIGIN/lib'
Compiling 21 files (.ex)
Generated exla app
Interactive Elixir (1.14.0) - press Ctrl+C to exit (type h() ENTER for help)
iex(1)> AxonBasic.run()

20:16:15.090 [info] could not open file to read NUMA node: /sys/bus/pci/devices/0000:01:00.0/numa_node
Your kernel may have been built without NUMA support.
 
20:16:15.090 [info] XLA service 0x7f33e00187b0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:

20:16:15.090 [info]   StreamExecutor device (0): NVIDIA GeForce RTX 3060, Compute Capability 8.6

20:16:15.090 [info] Using BFC allocator.

20:16:15.090 [info] XLA backend allocating 10641368678 bytes on device 0 for BFCAllocator.
 
20:16:15.445 [info] Start cannot spawn child process: No such file or directory
Epoch: 0, Batch: 10000, loss: 0.0101714
Epoch: 1, Batch: 10000, loss: 0.0067475
Epoch: 2, Batch: 10000, loss: NaN65432
Epoch: 3, Batch: 10000, loss: NaN
Epoch: 4, Batch: 10000, loss: NaN
Epoch: 5, Batch: 10000, loss: NaN
Epoch: 6, Batch: 10000, loss: NaN
Epoch: 7, Batch: 10000, loss: NaN
Epoch: 8, Batch: 10000, loss: NaN
Epoch: 9, Batch: 10000, loss: NaN
```

# おわりに

Cudaのドライバは11.8なのに、Ubuntu側のcudaパッケージは11.1で大丈夫なのか？という疑問はあるんですが、Axonを使った簡単なプログラムは動作しました。

[XLAのドキュメント](https://github.com/elixir-nx/xla)にXLA_BUILD=trueにすれば、パッケージのビルドができるので、ubuntu側のcudaパッケージを最新にして、ビルドすればいいはずなんですが、私の環境では、コンパイル時にエラーがでてうまくいっていません。別の記事に記載しようとおもいます。

cuda11.8用のXLAドライバービルド何処かに公開されてないかなぁ。

参考
https://docs.nvidia.com/cuda/wsl-user-guide/index.html#installing-wsl2


2022/10/16
config.exsの記述を修正しました。
default_defn_optionsを指定しないとdefnが高速化されてませんでした。