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

自己紹介 Twitter/Qiita @GeekMasahiro 会社員　自社サービスのソフトウエア開発、運営（ AWS, Python, vue.js等） 趣味　Elixir、電子工作全般、ソーラー発電、 3DPrinter Elixirを2022年7月に出会って、趣味に追加。楽しいので本業にも活用したい。 2

今回の目標 目標 ● NxにEXLAを導入してGPUで演算を行えるようにする 環境 ● nvidiaのGPU（RTX3060)を使用 ● Windows11のWSL環境で使用する 3

GPU購入 NxでGPUを使ってみたい！ Ethereumのマイニングが終わって価格 も戻ってきた 買いしました。 4

GPU取り付け 大きさ的に、内蔵できなかったので、 PCIeの延長 ケーブルを購入して外置き GPU用の電源も購入 PCの電源と連動する仕組みは電子工作。 5

ドライバーインストール(Windows側) https://developer.nvidia.com/cuda/wsl 6

WSL側の環境構築 nvidia社のドキュメントにWSL環境についての記述がある。この通りで動作する。 https://docs.nvidia.com/cuda/wsl-user-guide/index.html#getting-started-with-cuda-on-wsl CUDA Toolkit 11.8 Downloads（内容は同じ） https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distri bution=WSL-Ubuntu&target_version=2.0&target_type=deb_network 7

EXLAで使う場合の注意 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 #バージョン指定する 最新のバージョン（1.1.8だとうまく動作しない）ので、 11.1.0-1を指定 xlaのissue報告中。https://github.com/elixir-nx/xla/issues/23 8

動作確認 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" ], force: false, ) Nx.add(Nx.tensor([1]), Nx.tensor([1])) iexのプロンプトで実行 9

実行結果 cudaで作成されている 10

デモ ● Config.exsの働き ● cpu/cudaでの速度比較 ● GPUが動いている様子を観察 11

EXLAを使用する設定 import Config config :nx, :default_backend, EXLA.Backend config :nx, :default_defn_options, compiler: EXLA 12

XLA_TARGETの指定 export XLA_TARGET=cuda111 mix deps.clean --all rm -Rf ~/.cache/xla/ mix deps.get iex -S mix cpuを使う場合は、XLA_TARGET=cpu ※変更した場合、deps.cleanが必要。~/.cache/xla/を削除しないとEXLAのコンパイルエラーになる。 13

GPU動作状況の確認(nvidia-smi) 14

速度比較 200x200,2000x2000の単位行列の内積の演算時間を比較 条件 200x200 Exla使用しない 2.03sec cpu 0.0752sec rate 26倍 条件 2000x2000 cpu 12.4sec cuda 0.911us rate 13.6倍 15