はじめに
こんにちは，wottoと申します．普段はテンソルネットワークと量子コンピュータの繋がりについてあれこれ考えています．
近年のテンソルネットワーク法の急速な発展に伴い，量子回路を一種のテンソルネットワークとみなすことで，状態ベクトルをfullに保持するナイーブな方法よりも遥かに高速に量子回路を古典的にシミュレートできることができることが分かってきました．その結果，例えばGoogleが10000年かかると主張したタスクは実はスパコンを用いてわずか304秒で実行できることが実証され1，その成果は2021年度のゴードン・ベル賞に輝きました23．
それら画期的な研究と同時並行的に，テンソルネットワークベースの量子回路シミュレータライブラリがここ1,2年で多く開発されてきました4．特にNVIDIAのcuQuantumやJetなどGPUをアクセラレータに利用したライブラリが多く開発され，既存のシミュレータを凌駕する性能が報告されています5．
本記事では，GPUを用いた簡単なテンソルネットワークの計算について，その簡単な実装と実行結果をお見せしたいと思います．
cuTENSORを用いたテンソル積計算
といっても，テンソルネットワークの計算は基本的に縮約順さえ決めてしまえばあとはひたすらテンソル積（行列積）を計算するだけですので6，今回の実装ではそこまで工夫した点はありません．テンソルネットワークの基礎や，量子回路のシミュレーションをどのようにテンソルネットワークの計算に落とし込むことができるかについては，例えば去年の私のAdvent calendarの記事をご参照ください．
テンソル積の計算ですが，今回はcuTENSORライブラリを用いたいと思います．これはテンソル積に特化したNVIDIA CUDA-Xのひとつで，V100やA100などcompute capabilityが6.0以上のGPUで用いることができます．また，A100に搭載されているTensor coreを効率的に利用することができるらしいです．
cuTENSORのダウンロードはこのページを参照してください．cuTENSORの概念的な説明はこのスライドに載っています．
cuTENSORを生で書いてもいいんですが，CuPyのほうでラッパーの実装がありましたので，簡単のために今回はこちらを使いたいと思います．
cuTENSORによるテンソル計算では，具体的には以下のような計算を実行することができます．
$$
D_{m,u,n,v}=αA_{m,h,k,n}B_{u,k,v,h}+βC_{m,u,n,v}
$$
indexの数や種類はこちら側で変えることが出来ます．
以下のテンソル積を計算する箇所でcuTENSORを用いています．フルのコードはここをご覧ください．
# tensor descriptorの作成 テンソルを表すndarrayを渡す
desc_l = cutensor.create_tensor_descriptor(self.left.tensor)
desc_r = cutensor.create_tensor_descriptor(self.right.tensor)
c = cp.zeros([dim for i in range(len(self.index))])
desc_c = cutensor.create_tensor_descriptor(c)
# modeの作成 indexを表すint/strのtupleを渡す
mode_l = cutensor.create_mode(*self.left.index)
mode_r = cutensor.create_mode(*self.right.index)
mode_c = cutensor.create_mode(*self.index)
# cuTENSORを使うとき
self.tensor = cutensor.contraction(1.0, self.left.tensor, desc_l, mode_l, self.right.tensor, desc_r, mode_r, 1.0, c, desc_c, mode_c)
# cuTENSORを使わないとき
# self.tensor = cp.einsum(self.cont_str, self.left.tensor, self.right.tensor, dtype=np.float64)
結果
簡単のため，今回は量子回路のシミュレーションではなく，ランダムに結合している10ノードのテンソルネットワークの計算を行います．各枝の次元は8です．
上のようなテンソルネットワークについて，まず最適な縮約順を見つけ，それからテンソル縮約を順に行っていきます．
シードを変えて100回実行し，そのテンソル積のみの平均計算時間と合計でかかった時間を測定しました．比較として，cuTENSOR，cuBLAS，シングルスレッドのnumpyの３つの状況で実行しました．用いたGPUはA100です．
| cuTENSOR | cuBLAS | numpy | |
|---|---|---|---|
| 平均的なテンソル縮約時間 | 0.0302[s] | 0.0420[s] | 6.657[s] |
| トータルのプログラム実行時間 | 21.13[s] | 21.83[s] | 681.7[s] |
シングルスレッドのnumpyよりはGPUを用いたほうが遥かに高速に計算できていることがわかります．
ただし，cuTENSORを用いたときのメリットはあまり見えない結果となりました．テンソル積の計算で最も律速されるのはindexの入れ替え部分ですので，その辺りをちゃんとこちら側で制御しないとナイーブな実装ではcuTENSORの性能をうまく引き出せないのかもしれません．
おわりに
近年のDNN+GPUの大成功により科学技術計算，計算物理学においてもスパコンとGPUを組み合わせる動きが加速しています7．ただしGPUは1つあたりのメモリが小さいので，問題を小分けにしてGPUに投げるスケジューラをうまく書かないと高速化は達成できません．その辺りは私のような（スパコン・GPUどっちも）素人にとっては非常につらいところなのですが，特にテンソルネットワーク法は本質的な計算が全てテンソル積で完結しており，さらに問題の分割・並列化も容易ですのでGPUアクセラレータの恩恵を受けやすいタスクであるとも言えます．
またこれから色々勉強していく予定ですので，何かアドバイスやコメント等あれば頂けると幸いです．
-
このニュースにより「やっぱり量子コンピュータなんて使い物にならなかったんだ！」というインフォデミックがインターネットで発生していました． ↩
-
とはいえ，根本的なアルゴリズム自体は２年前ほどに既に考えられており，彼らの結果はむしろスーパーコンピュータで並列実行するための様々なチューニングに起因します．あまり詳しくないのですが，公式ページを見る限り，ゴードン・ベル賞自体は数理的手法を評価するというよりむしろ実際の問題に適用したその結果自体が評価指標の一つらしいので，「Googleの量子超越性を打ち破った」というインパクトが高く評価されたということでしょうか． ↩
-
量子コンピュータの古典シミュレーション，あるいは量子系の古典シミュレーションはそれ自体が非常に大事なタスクです．もちろん先の例のような「量子超越タスク」の量子vs古典というアカデミックな意味でも面白いですが，たとえばこれから先量子コンピュータの量子ビット数が多くなっていったときに量子コンピュータがどのような振る舞いを見せるのか？というのを古典シミュレーションによりある程度見積もりを立てることができます．そしてその見積もりは，古典シミュレート可能な量子ビット数が増える，あるいはかける量子ゲート数が増えれば増えるほど正確になります． ↩
-
2020年度の未踏ターゲット事業でも似たようなPJがありました． ↩
-
例えばcuQuantumは，先程のGoogleのSycamoreのランダム量子回路サンプリングタスクについて，わずか10分で1サンプルを生成できると主張しています．ただしこれについてはあまり詳細がない（or見つけられてない）のでなんともいえません． ↩
-
本当は，メモリに載りきるようにうまく並列処理などする必要があります． ↩
-
最近の大学とかが運営しているスパコンにはよくGPUが載ってたりします． ↩