NVIDIA Jetsonで材料シミュレーション


はじめに

過去にRaspberry Piでの材料シミュレーションに挑戦し、動作することは確認できましたが、そのパフォーマンスは実用的ではありませんでした。その後、性能が向上した新機種が登場していますが、劇的な高速化は期待できそうにありません。

懲りずに、Jetsonで再挑戦します。数値計算に利用するためのGPU (GPGPU)を搭載しているので、対応しているソフトウェアでは高性能が期待できます。最近LAMMPSを利用する機会が増えているので、同ソフトで高い性能を発揮できれば良しとします。最上位機種Jetson AGX Xavierの開発者キットを購入しました。Raspberry Piとはお値段が一桁違うので、それなりの性能が得られないと困ります。

セットアップ

付属ドキュメントは親切とは言い難いです。開封後、何から手を付ければよいのか皆目見当がつかず、ネット検索でこちらもページにたどり着きました。

こちらを見てから、再度同梱物を調べると、茶色の紙が説明書であることが理解できました。絵に描かれている通り操作すれば概ね大丈夫です。

上記サイトと私の作業の違いは、

  • 「初期セットアップ」はしませんでした。
  • ダウンロード完了後の画面で、Jetson側の情報を入力しますが、Jetsonを直接操作して設定しました。ユーザー名とパスワードを設定し、それをホストPCに入力しました。

「開発者キット」の名前通り、普通にインストールすれば開発環境が整います。MPIまで用意されていることには驚かされました。一方、Fortranはありません。

LAMMPSのコンパイル

LAMMPSのコンパイルは、あっけないほど簡単です。

tar xvf lammps-stable.tar.gz
cd lammps-29Sep2021
mkdir build; cd build
cmake -C ../cmake/presets/most.cmake -C ../cmake/presets/nolib.cmake -D PKG_GPU=on -D GPU_API=cuda -D GPU_ARCH=sm_72 ../cmake/
cmake --build .

GPU_ARCHは下記サイトを参照しました。

実行時に-sf gpuオプションを付加すると、GPU利用して計算されます。tegrastatsコマンドでGPUが利用されていることを確認します。

遅くはありませんが特段速いとも言えず、お値段相応といったところでしょうか。メモリを32GBも搭載しているので、これを有効活用できる場合には良い選択肢になる可能性があります。

Psi4のコンパイル

Psi4は、オープンソースの量子化学計算ソフトウェアです。MPI並列に対応していないので大規模計算に適しませんが、Pythonから操作するためのインターフェースが備わっていることが魅力です。
公式サイトのドキュメントを参考に、ソースからコンパイルします。

頼れることはconda-forgeを頼ることにして、Miniforge$HOME/miniforge3にインストールしてから、以下のコマンド列でPsi4をインストールします。

conda install scipy matplotlib py-cpuinfo psutil rdkit=2020.3.6  # Psi4と組み合わせて使いたい
conda install gfortran=9 gxx=9 cmake numpy # FortranはDFT-D3用
conda install networkx pint pydantic msgpack-python mpfr eigen

# Psi4
git clone https://github.com/psi4/psi4.git
cd psi4
git checkout cd00f19  # ver. 1.4.1相当
cmake -S . -Bobjdir -DCMAKE_INSTALL_PREFIX=$HOME/miniforge3
cd objdir
export CPLUS_INCLUDE_PATH=$HOME/miniforge3/include
make -j 4
make install

# RESP:これはpure Python
git clone https://github.com/cdsgroup/resp.git
cd resp
python setup.py install --prefix $HOME/miniforge3/

# DFT-D3:外部コマンドとして利用
# https://www.chemie.uni-bonn.de/pctc/mulliken-center/software/dft-d3/dftd3.tgz
tar zxf dftd3.tgz
cd dftd3
vi Makefile  # Makefile編集(下記参照;-staticを外しました)
make
cp dftd3 $HOME/miniforge3/bin/

DFTD3のMakefile修正は以下の通り。

--- Makefile.org        2016-06-08 10:30:18.000000000 +0900
+++ Makefile    2021-10-01 18:46:23.707810294 +0900
@@ -21,7 +21,7 @@ ifeq ($(OSTYPE),LINUXL)
 # CC = gcc
 # LINKER = lf95
 #  LINKER = ifort -static
-  LINKER = gfortran -static
+  LINKER = gfortran #-static
   PREFLAG = -E -P
   CCFLAGS = -O -DLINUX
        FFLAGS= -O

実行時にはPYTHONPATH$HOME/miniforge3/libを追加してください。

Psi4の実行速度は「遅い」です。LAMMPSの実行を遅いと感じなかったのは、GPGPUが効果を発揮したおかげかも知れません。

NVIDIA HPC SDK

NVIDIAからはHigh Performance Computing用の開発環境としてNVIDIA HPC SDKが提供されています。執筆時点の最新版21.9をインストールします。

ディスク容量が非常に厳しいので、増設することをお勧めします。動作確認(遊び)程度であれば増設せずとも利用できますので、その場合は必要ないファイルをできるだけ削除してから作業してください。

ライセンスに同意してから、Bundled with the newest plus two previous CUDA versions (11.4, 11.0, 10.2)Linux Arm Server DEBを選んで表示されるコマンドを実行します。
インストール先は/opt/nvidia/hpc_sdkです。.bashrcなどで環境変数を設定してください。

export PATH=/opt/nvidia/hpc_sdk/Linux_aarch64/21.9/comm_libs/mpi/bin:/opt/nvidia/hpc_sdk/Linux_aarch64/21.9/compilers/bin:$PATH
export LD_LIBRARY_PATH=/opt/nvidia/hpc_sdk/Linux_aarch64/21.9/comm_libs/mpi/lib:/opt/nvidia/hpc_sdk/Linux_aarch64/21.9/compilers/lib:$LD_LIBRARY_PATH

Fortranのコンパイルコマンドはnvfortran(もしくはmpif90)、Cのコンパイルコマンドはnvcです。
数値演算ライブラリが/opt/nvidia/hpc_sdk/Linux_aarch64/21.9/math_libsに用意されている気配を感じますが、残念ながらCUDA 10.2用の実体はありません。BLAS, LAPACKは/opt/nvidia/hpc_sdk/Linux_aarch64/21.9/compilers/libにあります。

PHASE/0

第一原理バンド計算PHASE/0を使って、こちらと同じ計算課題で実行時間を比較しました。MPIのみ8並列(k点2並列、バンド4並列)です。

結果:6分半

「速い」とは思いませんが、GPGPUを利用していないことを考慮すると、(コンパイラが?)頑張っているのではないでしょうか。

まとめ

目的が適切に設定されていれば、役立ちます。

  • GPGPUに対応しているソフトウェアを使う。例)LAMMPS
  • GPGPUを含む開発環境を安価に構築する。

深層学習フレームワークを利用する際にも有益であることが期待できますので、もう少し遊んでみたいと思います。