Zenn
⛳

Barrett Reduction (2)

2022/12/11に公開
Python
NVIDIA
CUDA
tech

Barrett Reduction (2) - on GPU -

はじめに

前回は、簡単にBarrett Reductionについて記載しました。
今回は、Pythonで並列計算(GPU)を利用して実行する方法について備忘録を残します。

使用パッケージ

とりあえず実装

CuPyのRawKernelを用いて実装します。

import cupy as cp

BarrettReduction = cp.RawKernel(r'''
extern "C" __global__ 
void BarretReduction(unsigned long long a[], unsigned int q, unsigned int mu, int q_bit_length)  
{
    int local_tid = threadIdx.x;

    unsigned long long rx;
    rx = a[local_tid] >> (q_bit_length - 2);
    rx = rx * mu;
    rx >>= q_bit_length + 2;
    rx *= q;
    a[local_tid] -= rx;

    a[local_tid] -= q * (a[local_tid] >= q);
}
''', 'BarretReduction')

上で定義したカーネルをPythonから呼び出して実行します。
なお、今回はグリッド数を1、スレッド数を最大の1024として呼び出します。

import ctypes

# Parameter initialization
q = 536608769
psi = 284166
psiinv = 208001377
q_bit_length = int(q.bit_length())
mu = int(cp.floor(2**(2*q_bit_length) / q))

a_cpu = cp.ones(1024, dtype=cp.uint64)
a_cpu *= 1024
a_cpu *= 393388074
d_a = cp.copy(a_cpu)

BarrettReduction((1,), (1024,), (d_a, q, mu, q_bit_length))

print(d_a)
# => [372811026 372811026 372811026 ... 372811026 372811026 372811026]

3. 今後やる予定のこと

  • Barrett Reductionを用いるアルゴリズム全般の実装と速度検証

参考文献

Discussion