Google ColaboratoryでCirq + cuStateVecの環境構築

Google Colaboratoryで量子コンピューティングのシミュレーションを行うため、Accelerating Quantum Circuit Simulation with NVIDIA cuStateVecを参考にCirq + cuStateVecの環境構築をしてみました。
手順通りだと動かなかったのでメモ。

インストール

上記のブログの通りcuquantum-linux-x86_64-0.1.0.30を選んでダウンロードし展開。
最新のバージョンを使っても動かないので注意。

! wget https://developer.download.nvidia.com/compute/cuquantum/redist/cuquantum/linux-x86_64/cuquantum-linux-x86_64-0.1.0.30-archive.tar.xz
! tar -xf cuquantum-linux-x86_64-0.1.0.30-archive.tar.xz
! export CUQUANTUM_ROOT=`pwd`/cuquantum-linux-x86_64-0.1.0.30-archive && ln -sf $CUQUANTUM_ROOT/lib $CUQUANTUM_ROOT/lib64

qsimをcheckoutしてソースからインストール。
同時にcirqもバージョン指定でインストール。
ビルドには数分程度かかりました。
cirqはこのバージョンでないと互換性がないためエラーになります。
pipに古いバージョンの依存関係を解決をさせるため、qsimとcirqは同時にインストールする必要があります。

! git clone https://github.com/quantumlib/qsim.git
! cd qsim && git checkout v0.11.1
! export CUQUANTUM_ROOT=`pwd`/cuquantum-linux-x86_64-0.1.0.30-archive && cd qsim && pip install . cirq==0.13.1

動作検証

GPUが使えることを確認。
detect_gpu()が0、detect_custatevec()が1を返せば成功。それ以外を返した場合はGPUが使えていません。(enum)

import qsimcirq

print("detect_gpu:", qsimcirq.qsim_decide.detect_gpu())
print("detect_custatevec:", qsimcirq.qsim_decide.detect_custatevec())

## 以下実行結果
# detect_gpu: 0
# detect_custatevec: 1

このissueを参考に実行時間を計測してみます。
リソースモニターでシステム RAMはほとんど変化せず、GPU RAMは上昇することを確認。
ここではGPUはT4を選択しています。

import cirq
import qsimcirq
import timeit

def load_test_with_gpu(num_gpus = 1, depth = 2, num_qubits = 30):
	circuit = cirq.testing.random_circuit(
		qubits = num_qubits,
		n_moments = depth,
		op_density = 1.0,
		random_state = 1)
	num_gates = len(list(circuit.all_operations()))
	options = qsimcirq.QSimOptions(use_gpu=True, gpu_mode=num_gpus)
	qsim_simulator = qsimcirq.QSimSimulator(options)
	result = qsim_simulator.simulate(circuit)
	print (f"DONE with qubits: {num_qubits} \t gates: {num_gates} \t depth: {depth} \t")

timeit.timeit('load_test_with_gpu()', globals=globals(), number=1)

## 以下実行結果
# DONE with qubits: 30 	 gates: 40 	 depth: 2 	
# 8.473605410000118

CPUでも実行してみます。
リソースモニターでシステム RAMが上昇してもGPU RAMはほとんど変化しないことを確認します。
GPUに比べて実行時間が約3.2倍になっています。

def load_test_with_cpu(depth = 2, num_qubits = 30):
	circuit = cirq.testing.random_circuit(
		qubits = num_qubits,
		n_moments = depth,
		op_density = 1.0,
		random_state = 1)
	num_gates = len(list(circuit.all_operations()))
	options = qsimcirq.QSimOptions(use_gpu=False)
	qsim_simulator = qsimcirq.QSimSimulator(options)
	result = qsim_simulator.simulate(circuit)
	print (f"DONE with qubits: {num_qubits} \t gates: {num_gates} \t depth: {depth} \t")

timeit.timeit('load_test_with_cpu()', globals=globals(), number=1)

## 以下実行結果
# DONE with qubits: 30 	 gates: 40 	 depth: 2 	
# 26.70925836099991