Python3 - numba パッケージの使い方

ホーム TeqStock.tokyo について 技術置き場 運営ブログ お問い合わせ
Python3 には CUDA 等の GPGPU を行うための numba パッケージが用意されています。
この記事では、numba パッケージの使い方について簡単に説明します。
CUDA のインストールが済んでいない場合は、こちらの記事を先にご覧ください。

numba のインストール

インストールは pip コマンドで行うことができます。
Anaconda3 を利用している場合は最初から組み込まれています。

環境変数の設定

利用前に CUDA の一部のライブラリの場所を環境変数に設定する必要があります。 
NUMBAPRO_NVVM=/usr/local/cuda/nvvm/lib64/libnvvm.so.3.2.0
NUMBAPRO_LIBDEVICE=/usr/local/cuda/nvvm/libdevice
.bashrc の最後に設定するようにすると良いと思います。

サンプルプログラム

以下は、２次元正方行列の転置行列を求めるプログラムです。

# -*- coding: utf-8 -*-
import numpy as np
from numba import cuda

# 転置行列を求める(デバイス側実装)
@cuda.jit
def vecT(b, a, n):
  x = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
  y = cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
  if x < n and y < n:
    b[y][x] = a[x][y]

# CUDA を使って転置行列を求める
def T(a):
  n = len(a)
  b = np.zeros((n, n), dtype=np.float32)
  threadN = (4, 4)
  blockN = ((threadN[0] - 1 + n) // threadN[0], (threadN[1] - 1 + n) // threadN[1])
  vecT[blockN, threadN](b, a, n)
  return b
 
if __name__ == "__main__":
  a = np.arange(0, 16, dtype=np.float32)
  a = a.reshape((4, 4))
  b = T(a)
  print(b)
cuda ライブラリを利用するには、numba パッケージから cuda を import します。
CUDA のエントリ関数(C++ では __global__ 修飾子を設定する関数)は @cuda.jit とアノテーション表記します。
もし、デバイス関数を定義したい場合は、@cuda.jit(device=True) のように表記します。

CUDA 関数に多次元配列を指定する場合、その順序は C-like つまり、列を先に指定する必要があります。
サンプルプログラムでは b[y][x] で要素 (x, y) にアクセスできたことになります。

CUDA エントリを呼び出すには block と thread を計算します。thread の最大値はアーキテクチャにより異なります。(最近のアーキテクチャは 8×8 でも可能)
正方行列は一辺が n 個あるため、block 数は n を thread 数で割って切り上げた値となります。
CUDA エントリの呼び出しは <関数名>[<block数>, <thread数>](パラメータ, ...) とします。
ブロック数とスレッド数の指定を忘れてもエラーにはなりませんが、エントリが呼び出されないので気をつけましょう。

block と thread は１次元しか使わないのであればタプルに詰める必要はなく、スカラを指定することができます。

Numpy の浮動小数点はデフォルトで float64 を使いますが、GPU には float32 を使う機会が多いため、サンプルプログラムでは float32 を指定しました。
さて、CUDA の効果ですが、サンプルプログラム程度では CUDA の方が圧倒的に遅いです。
やはり、行列サイズが一万くらいあり、ループも多重になっているシーンで真価を発揮してくれるようです。

Copyright (c) 2017-2019 by TeqStock.tokyo