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

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

環境構築

GPUを購入する

PCに取り付ける

Windows用Driverインストール

Run Linux GUI apps on the Windows Subsystem for Linux に従ってvGPUのドライバ（Windows側）をインストールする。

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

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

WSL2の環境構築(Not docker)

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)

デバイス認識確認

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に記載するなど。

.bashrc

export XLA_TARGET=cuda111

depsに、exlaを追加

mix.exs

  defp deps do
    [
      {:axon, "~> 0.2.0"}, # axonを使う場合
      {:exla, "~> 0.3.0"},
      {:nx, "~> 0.3.0"},
    ]
  end

configに追加

config/config.exs

import Config
config :nx, :default_backend, EXLA.Backend
config :nx, :default_defn_options, compiler: EXLA

プログラムの対応

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

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