Cuda C を用いた並列リダクションの実装と高速化

最近, cuda C を使う機会が多くて, リファレンスを探すのに苦労したので, そのメモです。
この記事は,

という流れです。

atomicAdd 関数の自作

複数のスレッドである変数を読み取り, 四則演算などをしてその値を書き込むときに使用します。
https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
倍精度浮動小数点数同士の和も, Compute Capability 6.0 以上ではatomicAdd が使えるらしい。そうでない場合は以下のように組む。

__device__ double atomicAdd_double(double* address, double val)
{
    unsigned long long int* address_as_ull =(unsigned long long int*)address;
    unsigned long long int old = *address_as_ull, assumed;

    do {
        assumed = old;
        old = atomicCAS(address_as_ull, assumed,
                        __double_as_longlong(val + __longlong_as_double(assumed)));
    } while (assumed != old);

    return __longlong_as_double(old);
}

並列リダクションの実装

https://devblogs.nvidia.com/faster-parallel-reductions-kepler/
をみた感じだと, CUB Library を使用するか warp+atomicAdd を使用すると良さそう。
version が古いと, double 型に対する shuffle_down 操作も書く必要があるので, まずはそのバージョン。

__device__ double __shfl_down_double(double var, unsigned int srcLane, int width = 32) {
	int2 a = *reinterpret_cast<int2*>(&var);
	a.x = __shfl_down_sync(0xffffffff, a.x, srcLane, width); // or a.x = __shfl_down(a.x, srcLane, width);
	a.y = __shfl_down_sync(0xffffffff, a.y, srcLane, width); // or a.y = __shfl_down(a.y, srcLane, width); 
	return *reinterpret_cast<double*>(&a);
}

__device__ double warpReduceSum_double(double val) {
	for (int offset = warpSize / 2; offset > 0; offset /= 2)
		val += __shfl_down_double(val, offset);
	return val;
}

__global__ void deviceReduceWarpAtomicKernel(double *in, double* out, int N) {
	double sum = double(0.0);
	for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += blockDim.x * gridDim.x) {
		sum += in[i];
	}
	sum = warpReduceSum_double(sum);
	if ((threadIdx.x & (warpSize - 1)) == 0)
		atomicAdd_double(out, sum);
}

deviceReduceWarpAtomicKernel のfor 文では, grid stride loop が用いられています。
CUDA Pro Tip: Write Flexible Kernels with Grid-Stride Loops | NVIDIA Developer Blog

Programming Guide :: CUDA Toolkit Documentation
を見てみると, __shfl_down_sync は int, unsigned int, long, unsigned long, long long, unsigned long long, float, doubleに対応しているらしい。また
Using CUDA Warp-Level Primitives | NVIDIA Developer Blog
warp に関する情報が記載されています。
CUDA9.0 以降だと,

__device__ double warpReduceSum_double(double val) {
	for (int offset = (warpSize >> 1); offset > 0; offset >>= 1)
        val += __shfl_down_sync(0xffffffff, val, offset);
	return val;
}

__global__ void deviceReduceWarpAtomicKernel(double *in, double* out, int N) {
	double sum = double(0.0);
	for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += blockDim.x * gridDim.x) {
		sum += in[i];
	}
	sum = warpReduceSum_double(sum);
	if ((threadIdx.x & (warpSize - 1)) == 0)
		atomicAdd_double(out, sum);
}

で, OK。

ベクトル化による高速化(あまり効果はなかった)

さらに, 先ほどのコードをベクトル化することにより高速化してみる。
CUDA Pro Tip: Increase Performance with Vectorized Memory Access
が参考になります。

__global__ void deviceReduceWarpAtomicVectorized2Kernel(double *in, double* out, unsigned int N) {
    double sum = double(0.0);
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int half_N = N >> 1;
    double2 tmp;
    for (unsigned int i = idx; i < half_N; i += blockDim.x * gridDim.x) {
        tmp = reinterpret_cast<double2*>(in)[i];
        sum += tmp.x + tmp.y;
    }

    sum = warpReduceSum_double(sum);
    if ((threadIdx.x & (warpSize - 1)) == 0) atomicAdd_double(out, sum);
}

以下では, N = 2**26 として, 実行します。
nvprof を用いて大雑把に実行時間の比較:
deviceReduceWarpAtomicKernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024>>> 798.97
<<< 131072, 512>>> 746.29
<<< 262144, 256>>> 573.90
<<< 524288, 128>>> 577.57
<<< 131072, 256>>> 286.78
<<< 65536, 256>>> 143.56
<<< 32768, 256>>> 72.505
<<< 16384, 256>>> 37.565
<<< 8192, 256>>> 20.420
<<< 4096, 256>>> 12.038
<<< 2048, 256>>> 8.5510
<<< 1024, 256>>> 8.0172
<<< 512, 256>>> 7.6053
<<< 256, 256>>> 7.5142
<<< 128, 256>>> 7.4906

deviceReduceWarpAtomicVectorized2Kernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024 >>> 408.98
<<< 131072, 512 >>> 385.48
<<< 262144, 256 >>> 298.33
<<< 524288, 128 >>> 300.19
<<< 131072, 256 >>> 282.86
<<< 65536, 256 >>> 141.58
<<< 32768, 256 >>> 71.043
<<< 2048, 256 >>> 8.2606
<<< 1024, 256 >>> 7.5353
<<< 512, 256 >>> 7.4422
<<< 256, 256 >>> 7.4634
<<< 128, 256 >>> 7.4942

同様にして,

__global__ void deviceReduceWarpAtomicVectorized4Kernel(double *in, double* out, unsigned int N) {
    double sum = double(0.0);
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int loop_N = N >> 2;
    double4 tmp;
    for (unsigned int i = idx; i < loop_N; i += blockDim.x * gridDim.x) {
        tmp = reinterpret_cast<double4*>(in)[i];
        sum += tmp.x + tmp.y + tmp.z + tmp.w;
    }

    sum = warpReduceSum_double(sum);
    if ((threadIdx.x & (warpSize - 1)) == 0) atomicAdd_double(out, sum);
}

deviceReduceWarpAtomicVectorized4Kernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024>>> 213.49
<<< 131072, 512>>> 180.72
<<< 262144, 256>>> 166.55
<<< 524288, 128>>> 167.64
<<< 65536, 256>>> 142.91
<<< 2048, 256>>> 8.2487
<<< 1024, 256>>> 7.7647
<<< 512, 256>>> 7.9956

ベクトル化はやりすぎても良くないみたいで, 効果はいまいちみたいです… ベクトル化を使う場合にdouble2 までというのは事実っぽいですね。それ以上に, ブロック当たりの thread 数やグリッド当たりのブロック数のチューニングをする方が, 性能がかなり変わりますね!

複素数の場合

同様のことが複素数に対してもできます。

__global__ void deviceReduceWarpAtomicKernel(cuDoubleComplex *in, cuDoubleComplex* out, unsigned int N) {
	cuDoubleComplex sum = make_cuDoubleComplex(0.0, 0.0);
	for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += blockDim.x * gridDim.x) {
		sum = cuCadd(sum, in[i]);
	}
	sum.x = warpReduceSum_double(sum.x);
	sum.y = warpReduceSum_double(sum.y);
	if ((threadIdx.x & (warpSize - 1)) == 0){
		atomicAdd_double(&out[0].x, sum.x);
		atomicAdd_double(&out[0].y, sum.y);
    }
}

__global__ void deviceReduceWarpAtomicVectorized2Kernel(cuDoubleComplex *in, cuDoubleComplex* out, unsigned int N) {
	cuDoubleComplex sum = make_cuDoubleComplex(0.0, 0.0);
    double4 tmp;
    unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int half_N = N >> 1;
	for (unsigned int i = idx; i < half_N; i += blockDim.x * gridDim.x) {
        tmp = reinterpret_cast<double4*>(in)[i];
		sum.x += tmp.x + tmp.z;
        sum.y += tmp.y + tmp.w;
	}
	sum.x = warpReduceSum_double(sum.x);
	sum.y = warpReduceSum_double(sum.y);
	if ((threadIdx.x & (warpSize - 1)) == 0){
		atomicAdd_double(&out[0].x, sum.x);
		atomicAdd_double(&out[0].y, sum.y);
    }
}

複素数の場合の実行時間の比較:
deviceReduceWarpAtomicKernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024>>> 831.83
<<< 131072, 512>>> 796.84
<<< 262144, 256>>> 574.31
<<< 524288, 128>>> 577.55
<<< 256, 256>>> 15.069
<<< 128, 256>>> 14.850

deviceReduceWarpAtomicVectorized2Kernel

<<< block per grid, threads per block >>> execution time (ms)
<<< 65536, 1024 >>> 438.12
<<< 131072, 512 >>> 412.84
<<< 262144, 256 >>> 319.97
<<< 524288, 128 >>> 321.76
<<< 131072, 256 >>> 290.40
<<< 4096, 256 >>> 17.290
<<< 2048, 256 >>> 15.544
<<< 1024, 256 >>> 15.709
<<< 256, 256 >>> 15.889
<<< 128, 256 >>> 16.083

ブロックあたりのスレッド数は 256 で良さそうなので, 以下ではスレッド数を256 とします。

高速化の結果

さらに, 同様の方針で自作をして実行時間を計測したものをまとめると以下の表のようになります。

関数名 実数の場合の実行時間 (ms) 複素数の場合の実行時間 (ms)
deviceReduceWarpAtomicKernel (事実上 grid stride なし) 798.97 831.83
deviceReduceWarpAtomicKernel 7.4906 14.850
deviceReduceWarpAtomicVectorized2Kernel 7.4422 15.544
deviceReduceWarpAtomicVectorized4Kernel 7.7647 24.040

並列リダクションを行う場合, atomic 操作がボトルネックになりやすいので, それを減らすようにすると速くなるというぐらいです。

Quantum Kitchen Sinks

[1806.08321] Quantum Kitchen Sinks: An algorithm for machine learning on near-term quantum computers
が気になったので、そのメモ。

どんなもの?
Quantum Kitchen Sinks(QKS) は、一連の量子回路からランダムサンプリングし、各々の回路を用いて入力データを測定されたビット列へ非線形変換を実現するものである。このアルゴリズムによって量子回路のパラメータの最適化にかかるコストを排除することができる!!そのあと、連結された結果は、(量子ではない通常の)機械学習アルゴリズムを用いて処理する(量子回路の効果が分かるように、単純なものを用いる)。
f:id:puyokw:20180626230407p:plain:w400

技術や手法のキモはどこ?
Quantum Kitchen Sinks(QKS) は、Random Kitchen Sinks(RKS) にinspire されたもので、RKS では例えばcos 関数が使用されているが、これは単一の量子ビットにRabi 回転(下図では, RX) を適用することでそのような変換ができる。そのときのパラメータはランダムに決める。

f:id:puyokw:20180626233156p:plain:w400
上図の(a) とか(c) はうまくいくが、(b) はダメらしい。

入力データは実数だけど、どうやって量子回路に入れるの??
前処理:
 \Theta_{i,e} = \Omega_e u_i + \beta_e
ただし,  \Omega_e \mathcal N(0, \sigma^2) に,  \beta_e \mathcal U(0, 2\pi) に従う。
この \theta RX(\theta) に入れることになります。

メインの量子回路の処理を行う。

後処理:観測を行って, 出力される0, 1 を全エピソード数E 個分つなげていく(ここの部分は, 複数回観測した合計や平均などでも可) 。パラメータ付き量子回路のパラメータ数が4 個なら, 各トレーニングデータは長さ4E のベクトルになる。
E は大きい方が良い。

どうやって有効だと検証した?
(1) "picture frames" の二値分類
|- (a) logistic regression (scikit-learn) (~50%)
|- (b) QKS + logistic regression (> 99.9%)

(2) MNIST の数字の3 と5 の二値分類
|- (a) logistic regression (95.9%)
|- (b) QKS + logistic regression (98.6%)

ロジスティク回帰のみと量子回路+ロジスティク回帰を比較していて、QKS の段階で量子回路をたくさん作ることで精度を向上させることができる。
f:id:puyokw:20180626231513p:plain

もう一つの戦略として, qubit 数を増やすのもあるが, こちらでは途中から過学習している。
f:id:puyokw:20180626231538p:plain:w400
これは4 qubit の場合の例で, qubit 数を増やしたら精度が上がりそうに思うのですが…
f:id:puyokw:20180627012940p:plain:w200

議論はある?
・CNOT を含むある特定のネットワークが準最適である可能性がある。
・MNIST のデータセットでは、大きな回路のパラメータをチューニングするには不十分な大きさである。

今回試されていた回路は例えば以下のようなRX とCNOT を組み合わせたものでした。
例えば、4 qubit の場合だと以下のような構成。
RX(%x0) 0
RX(%x1) 1
RX(%x2) 2
RX(%x3) 3
CNOT 0 2
CNOT 1 3
CNOT 0 1
CNOT 2 3

感想:
量子回路の設計方針が分かりやすい!


TO DO
・再現実装をしてみる(まだ雰囲気程度)

from pyquil.quil import Program
from pyquil.api import QVMConnection
from pyquil.gates import *
from sklearn import datasets
from sklearn.linear_model import LogisticRegression
import numpy as np
import math
from tqdm import trange
from sklearn.metrics import log_loss, accuracy_score
from sklearn.cross_validation import StratifiedKFold
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import pandas as pd

def rand():
	return np.random.rand()

def rand_u():
	return np.random.ranf()*2-1

def rand_n():
	return np.random.normal(0,1)

def make_pictureframe(num_data):
	data=np.zeros((num_data,2))
	teacher=np.zeros(num_data).astype(np.int32)
	for i in range(num_data):
		if rand()<=0.5:
			teacher[i]=0
		else:
			teacher[i]=1
	for i in range(num_data):
		tmp = rand()
		if tmp <= 0.25:
			data[i][0] = 1.0*rand_u() + rand_u()*0.1
			if rand() <= 0.5:
				data[i][1] = -1.0 + rand_u()*0.1
			else:
				data[i][1] = 1.0 + rand_u()*0.1
		elif tmp <= 0.5:
			data[i][0] = -1.0*rand_u() + rand_u()*0.1
			if rand() <= 0.5:
				data[i][1] = -1.0 + rand_u()*0.1
			else:
				data[i][1] = 1.0 + rand_u()*0.1
		elif tmp <= 0.75:
			data[i][1] = -1.0*rand_u() + rand_u()*0.1
			if rand() <= 0.5:
				data[i][0] = -1.0 + rand_u()*0.1
			else:
				data[i][0] = 1.0 + rand_u()*0.1
		else:
			data[i][1] = 1.0*rand_u() + rand_u()*0.1
			if rand() <= 0.5:
				data[i][0] = -1.0 + rand_u()*0.1
			else:
				data[i][0] = 1.0 + rand_u()*0.1
	for i in range(num_data):
		if teacher[i] == 1:
			data[i][0] *= 2
			data[i][1] *= 2
	return [data, teacher]


#np.random.seed(1)
data, target = make_pictureframe(100)

num_param=2
num_data=data.shape[0]
p=data.shape[1]
q=num_param # number of qubit
n_qubits=num_param # number of qubit
r=int(p/q)

num_episode=200
sigma=1

#np.random.normal(0, sigma)

Omega=[]
beta=[]
for e in range(num_episode):
	Omega.append( (np.concatenate( (np.random.randn(r, q), np.zeros((p-r,q))), axis=0)).T )
	beta.append( np.random.rand(q)*2*math.pi)

Omega=np.array(Omega)
beta=np.array(beta)

data_transform=[]
for i in trange(num_data):
	tmp=[]
	for e in trange(num_episode):
		u=data[i]
		theta = np.dot(Omega[e], u) + beta[e]
		qvm = QVMConnection()
		p1=Program()
		for j in range(q):
			p1.inst(RX(theta[j],j))
		
		p1.inst(CNOT(0, 1))
		p1.inst(MEASURE(0,0))
		p1.inst(MEASURE(1,1))
		tmp.append(np.array(qvm.run(p1, [j for j in range(q)], 10000)).mean(axis=0))
		
	tmp=[item for sublist in tmp for item in sublist]
	data_transform.append(tmp)

data_transform = pd.DataFrame(data_transform)
data_transform.columns=["col_"+str(i) for i in range(data_transform.shape[1])]
#data_transform.to_csv("trained.csv", index=False)

#train = pd.read_csv("trained.csv")
train=data_transform
train = np.array(train)

skf=StratifiedKFold(target, n_folds=5, shuffle=True)
for tr, te in skf:
	X_train, X_valid, y_train, y_valid = train[tr], train[te], target[tr], target[te]

lr = LogisticRegression()
lr.fit(X_train, y_train)

pred=lr.predict_proba(X_valid)
test_logloss=log_loss(y_valid, pred)
test_acc=lr.score(X_valid, y_valid)
pred=lr.predict_proba(X_train)
train_logloss=log_loss(y_train, pred)
train_acc=lr.score(X_train, y_train)
print("train acc: {}, test acc: {}".format(train_acc, test_acc))
print("train logloss: {}, test logloss {}".format(train_logloss, test_logloss))

pyquil は遅いので, 他のシミュレータで試した感じだと85% 程度しか精度でない....
ノイズがあった方が良いのかな??

Rigetti のpyquil のお試し

IBMQ 以外にも量子プログラミングができるプラットフォームとして Rigetti Forest があります。
Rigetti
Rigetti Forest のシミュレータは, メールアドレスを登録することで使用できます。実機での実行するためには, アカウントのアップグレードが必要なため, ここではシミュレーションの方を試します。python のpyquil を使用します。
GitHub - rigetticomputing/pyquil: A Python library for quantum programming using Quil.

pip install pyquil

でインストールできます。

メールアドレスを登録すると, key とuser_id が貰えますので, それを以下のような内容の .pyquil_config ファイルを作成します。

[Rigetti Forest]
key: <Rigetti Forest API key>
user_id: <Rigetti User ID>

Installation and Getting Started — pyQuil 2.0.0.dev0 documentation
の注意書きにあるように, Linux では /users/username に, Mac では /Users/Username に, Windows では C:\Users\Username にこのファイルを配置します。

以下では, 回路の作成, 回路の図示を行います。

from pyquil.quil import Program
from pyquil.api import QVMConnection
from pyquil.gates import *
qvm = QVMConnection()
p1=Program()
p1.inst(H(0),
        CNOT(0,1),
        H(1),
        MEASURE(0,0),
        MEASURE(1,1)) # 回路を作成
print(p1)
# 出力
#H 0
#CNOT 0 1
#H 1
#MEASURE 0 [0]
#MEASURE 1 [1]

pyquil では柔軟に回路を作成することができます。

p1.inst(H(0),
        CNOT(0,1),
        H(1),
        MEASURE(0,0),
        MEASURE(1,1)) 
の部分は
p1.inst(H(0))
p1.inst(CNOT(0,1))
p1.inst(H(1))
p1.inst(MEASURE(0,0))
p1.inst(MEASURE(1,1)) 
と書いても同じです。

また複数の回路を結合することもできます。

q1=Program()
q1.inst(Z(0),I(0))
q1.inst(X(1))
print(q1)
# 出力
# Z 0
# I 0
# X 1
# q1 の後にp1 を結合する
q1.inst(p1)
print(q1)
# 出力結果
# Z 0
# I 0
# X 1
# H 0
# CNOT 0 1
# H 1
# MEASURE 0 [0]
# MEASURE 1 [1]

シミュレーションの実行は次のようにして行います。

result=qvm.run(p1, [0,1], 10)
print(result) 
# [[1, 0], [0, 1], [1, 0], [0, 0], [0, 0], [0, 1], [1, 0], [0, 0], [1, 0], [0, 1]]

試行回数を増やした版で再度実験して集計します。

result=qvm.run(p1, [0,1], 8000)
# 各サブリストの中をまとめる
L=[]
for sublist in result:
    L.append(''.join([str(l) for l in sublist]))

# pandas を用いて集計する
import pandas as pd 
import matplotlib.pyplot as plt
df=pd.DataFrame({"state": L})
df['state'].value_counts().plot.bar()
plt.show()

f:id:puyokw:20180408040019p:plain

from pyquil.latex.latex_generation import to_latex 
tmpdir='tmp/'
filename='circuit1'
dc=to_latex(q1.inst()) # q1 の回路をtex 形式に変換
with open(tmpdir+filename+'.tex','w') as f:
    f.write(dc) 

回路をtex 形式に変換するまではpyquil でできますが, そこからpdflatex を用いてpdf に変換します。それをさらにjpeg などに変換して読み込みます。

import os
os.system("pdflatex -output-directory {} {}".format(tmpdir, filename+".tex"))
os.system("pdftocairo -jpeg {}.pdf {}".format(tmpdir+filename, tmpdir+filename))

from PIL import Image 
img = Image.open(os.path.join('./'+tmpdir+filename+'-1.jpg'))
img.show()

f:id:puyokw:20180408035143j:plain
他にもsympy のqasm を利用することでも回路図を表示することもできます。

from sympy.physics.quantum.qasm import Qasm
from sympy.physics.quantum.gate import IdentityGate
import matplotlib.pyplot as plt

# Qasm では, identity gate が含まれていないため, 継承して追加します
class QasmEx(Qasm):
    def identity(self, arg):
        self.circuit.append(IdentityGate(self.index(arg)))

# 回路を図示
def VisualizeCircuit(circuit):
    STR=str(circuit).split('\n')[:-1]
    SSTR=[]
    for s in STR:
        s_tmp=s.split(' ')
        if s_tmp[0]=='MEASURE':
            s_tmp=s_tmp[:-1]
        s_tmp2=s_tmp[0]
        if len(s_tmp)>1:
            s_tmp2+=' '
            s_tmp2+=','.join(list(map(lambda x: 'q'+x,s_tmp[1::])))
        SSTR.append(s_tmp2)
    SSTR=list(map(lambda x: x.replace("MEASURE","measure").replace("I","identity"),SSTR))
    SSTR=["qubit "+"q"+str(i) for i in range(2)]+SSTR
    qprog1=QasmEx()
    for i in range(len(SSTR)):
        qprog1.add(SSTR[i])
    qprog1.plot()

VisualizeCircuit(q1.inst())
plt.show()

f:id:puyokw:20180409023857p:plain

pyquil の印象としては, 回路を書くのは便利そうな感じがしました。