7. 2次元配列の計算

7.1 2次元配列の計算

2次元配列の計算方法には各種あり、その方法によって計算時間に大きな差がでます。
配列演算を用いるとプログラムが簡単になり、計算時間も短縮されます。 さらに、プロードキャストという技術を用いるとプログラムがより簡単になります。
一方、Numbaを用いるとfor文のままで高速化することができます。 さらに、Numbaは配列演算と併用することもできます。
リスト7-1では同じ処理(c_ij=a_i+b_j) を15の方法で記述しています。

(1)～(5)が配列演算です。
配列全体をベクトル処理するには":"を用います。
(5)はブロードキャストの一例です。 (i, j)の2重ループにブロードキャストを適用するには、 iに関する配列aは a.reshape((n, 1)) とし、 jに関する配列bは b.reshape((1, n)) とします（n : 配列の大きさ）。

(6)～(15)ではNumbaを使用しています。
すべて、"@jit(cache=True, nopython=True, nogil=True, parallel=True)" を使用しています。
for文の処理には通常の"range"と並列処理をする"prange"があります。

リスト7-1 2次元配列の計算 (array2d.py)

# array2d.py
# test of 2D array

import numpy as np
import time
from numba import jit, prange, set_num_threads

# (1) i, j, loop
def fn1(a, b, c):
    n = len(a)
    for i in range(n):
        for j in range(n):
            c[i, j] = a[i] + b[j]

# (2) j, i, loop
def fn2(a, b, c):
    n = len(a)
    for j in range(n):
        for i in range(n):
            c[i, j] = a[i] + b[j]

# (3) j loop (array)
def fn3(a, b, c):
    n = len(a)
    for j in range(n):
        c[:, j] = a + b[j]

# (4) i loop (array)
def fn4(a, b, c):
    n = len(a)
    for i in range(n):
        c[i, :] = a[i] + b

# (5) broadcast
def fn5(a, b, c):
    n = len(a)
    c[:] = a.reshape((n, 1)) + b.reshape((1, n))

# (6) i, j, loop (JIT, range+range)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn6(a, b, c):
    n = len(a)
    for i in range(n):
        for j in range(n):
            c[i, j] = a[i] + b[j]

# (7) i, j, loop (JIT, prange+range)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn7(a, b, c):
    n = len(a)
    for i in prange(n):
        for j in range(n):
            c[i, j] = a[i] + b[j]

# (8) i, j, loop (JIT, range+prange)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn8(a, b, c):
    n = len(a)
    for i in range(n):
        for j in prange(n):
            c[i, j] = a[i] + b[j]

# (9) j, i, loop (JIT, range+range)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn9(a, b, c):
    n = len(a)
    for j in range(n):
        for i in range(n):
            c[i, j] = a[i] + b[j]

# (10) j, i, loop (JIT, prange+range)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn10(a, b, c):
    n = len(a)
    for j in prange(n):
        for i in range(n):
            c[i, j] = a[i] + b[j]

# (11) j, i, loop (JIT, range+prange)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn11(a, b, c):
    n = len(a)
    for j in range(n):
        for i in prange(n):
            c[i, j] = a[i] + b[j]

# (12) j loop (JIT, j : range, i : array)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn12(a, b, c):
    n = len(a)
    for j in range(n):
        c[:, j] = a + b[j]

# (13) j loop (JIT, j : prange, i : array)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn13(a, b, c):
    n = len(a)
    for j in prange(n):
        c[:, j] = a + b[j]

# (14) i loop (JIT, i : range, j : array)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn14(a, b, c):
    n = len(a)
    for i in range(n):
        c[i, :] = a[i] + b

# (15) i loop (JIT, i : prange, j : array)
@jit(cache=True, nopython=True, nogil=True, parallel=True)
def fn15(a, b, c):
    n = len(a)
    for i in prange(n):
        c[i, :] = a[i] + b

# parameters
N = 20000
fn = 1  # = 1...11
L = 1
thread = 8
dtype= 'f4'  # 'f4' or 'f8'

# setup
a = np.arange(N).astype(dtype)
b = np.arange(N).astype(dtype)
c = np.zeros((N, N), dtype)

# Numba threads
if fn >= 6:
    set_num_threads(thread)

t0 = time.time()

# calculation
for _ in range(L):
    if   fn == 1:
        fn1(a, b, c)
    elif fn == 2:
        fn2(a, b, c)
    elif fn == 3:
        fn3(a, b, c)
    elif fn == 4:
        fn4(a, b, c)
    elif fn == 5:
        fn5(a, b, c)
    elif fn == 6:
        fn6(a, b, c)
    elif fn == 7:
        fn7(a, b, c)
    elif fn == 8:
        fn8(a, b, c)
    elif fn == 9:
        fn9(a, b, c)
    elif fn == 10:
        fn10(a, b, c)
    elif fn == 11:
        fn11(a, b, c)
    elif fn == 12:
        fn12(a, b, c)
    elif fn == 13:
        fn13(a, b, c)
    elif fn == 14:
        fn14(a, b, c)
    elif fn == 15:
        fn15(a, b, c)

t1 = time.time()

# output
exact = N**2 * (N - 1)
print('(%d) N=%d(%d), %.2f[sec], %e, %e'
    % (fn, N, L, (t1 - t0) / L, np.sum(c), exact))

# free
a = None
b = None
c = None

7.2 計算時間

表7-1にそれぞれの手法の計算時間を示します。
計算時間の短いケースでは繰り返し回数Lを十分大きくとって総計算時間をLで割っています。
計算時間の短いケースに色をつけています。

表7-1から以下のことがわかります。

• (1)(2) すべてforループのときは非常に計算時間がかかる。 i,jの順序や単精度と倍精度の差も小さい。
• (3)(4) iのforループ(4)はjのforループ(3)より大幅に速い。 これはPythonが既定値では"row major"であることによる。
• (5) ブロードキャストを用いるとfor文をなくすことができるが、 必ずしも最速ではない。(4)と比べると遅い。
• (6)(7)(8) i,jの順の2重forループにNumbaを適用すると速くなる。 外側のループにprangeを適用したときが最も速くなり、 両方のループにrangeを適用したときがそれに続き、 内側のループにprangeを適用したときが最も遅い。 これは、並列計算の粒度は大きい方がよく、 また並列計算を多数回起動したときはそのオーバーヘッドが無視できないことによる。
• (9)(10)(11) j,iの順の2重forループは (6)(7)(8)のi,jの順の2重forループと比べて大幅に遅い。 (3)(4)と同じ理由である。
• (12)(13) 外側をjのforループ、内側をiの配列演算としたときは大幅に遅くなる。 prangeを使うと却って遅くなる場合がある。
• (14)(15) 外側をiのforループ、内側をjの配列演算としたときは速くなる。 ループについてはprangeを用いると並列計算の効果が大きい。
• (4)(6)(7)(15)のケースは最適化の効果が得られ速くなる。 特に(7)(15)のケースが同等の性能で最も速い。
• (4)のようにすべてを配列演算で記述するとプログラムが簡潔になり、 Pythonの標準機能なのでNumbaのような追加のパッケージは不要であるが、 速度はNumbaに及ばない。これはGILにより並列計算できないためと思われる。
• 最適化されたケースでは単精度は倍精度の約2倍速い。 それ以外のケースでは単精度と倍精度の速度差は小さい。

7.3 C版の計算時間

参考までに同じ計算をCで行った結果を調べます。
リスト7-2にソースコードを示します。

リスト7-2 C版のソースコード (array2d.c)

/*
array2d.c
test program of 2D array

$ cl.exe array2d.c /O2 /openmp
$ array2d.exe [fn N L threads]
*/
# include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <omp.h>

// (1) i, j loop, i : OpenMP
static void fn1(int n, float *a, float *b, float **c)
{
	int i;
#pragma omp parallel for
	for (i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			c[i][j] = a[i] + b[j];
		}
	}
}

// (2) i, j loop, j : OpenMP
static void fn2(int n, float *a, float *b, float **c)
{
	for (int i = 0; i < n; i++) {
		int j;
#pragma omp parallel for
		for (j = 0; j < n; j++) {
			c[i][j] = a[i] + b[j];
		}
	}
}

// (3) j, i loop, j : OpenMP
static void fn3(int n, float *a, float *b, float **c)
{
	int j;
#pragma omp parallel for
	for (j = 0; j < n; j++) {
		for (int i = 0; i < n; i++) {
			c[i][j] = a[i] + b[j];
		}
	}
}

// (4) j, i loop, i : OpenMP
static void fn4(int n, float *a, float *b, float **c)
{
	for (int j = 0; j < n; j++) {
		int i;
#pragma omp parallel for
		for (i = 0; i < n; i++) {
			c[i][j] = a[i] + b[j];
		}
	}
}

int main(int argc, char **argv)
{
	int fn      = 1;
	int N       = 2000;
	int L       = 1;
	int threads = 1;

	clock_t t0, t1;

	// arguments
	if (argc > 4) {
		fn      = atoi(argv[1]);
		N       = atoi(argv[2]);
		L       = atoi(argv[3]);
		threads = atoi(argv[4]);
	}

	// OpenMP
	omp_set_num_threads(threads);

	// alloc
	float *a = (float *)malloc(N * sizeof(float));
	float *b = (float *)malloc(N * sizeof(float));
	float **c = (float **)malloc(N * sizeof(float *));
	for (int i = 0; i < N; i++) {
		c[i] = (float *)malloc(N * sizeof(float));
	}

	// setup
	for (int i = 0; i < N; i++) {
		a[i] = i;
		b[i] = i;
	}

	t0 = clock();

	// calculation
	for (int l = 0; l < L; l++) {
		if      (fn == 1) {
			fn1(N, a, b, c);
		}
		else if (fn == 2) {
			fn2(N, a, b, c);
		}
		else if (fn == 3) {
			fn3(N, a, b, c);
		}
		else if (fn == 4) {
			fn4(N, a, b, c);
		}
	}

	t1 = clock();

	// check sum
	double sum = 0;
	for (int i = 0; i < N; i++) {
		for (int j = 0; j < N; j++) {
			sum += c[i][j];
		}
	}

	// output
	const double cpu = (double)(t1 - t0) / CLOCKS_PER_SEC / L;
	const double exact = (double)N * N * (N - 1);
	printf("(%d) N=%d(%d), %.2f[sec], %e, %e\n", fn, N, L, cpu, sum, exact);

	return 0;
}

表7-2にそれぞれの手法の計算時間を示します。

表7-2から以下のことがわかります。

• (1) このとき最も速くなる。並列計算の速度比は2倍以下である。
• (2) 1スレッドのときは(1)と同じはずであるが少し遅いのはコンパイラの最適化の問題か。 内側のループで並列計算しても(1)と比べて速度の劣化は小さい。
• (3) ループの順序を(j,i)に変えると大幅に遅くなる。 CではPythonと同じく"row major"であるためである。
• (4) 1スレッドのときは(3)と同じであるが、 内側のループを並列計算すると大幅に速くなる。ただし(1)(2)には及ばない。
• (1)の1スレッドの計算時間は表7-1の(4)と同じであり、 (1)の8スレッドの計算時間は表7-1の(7)(15)と同じである。 ほぼ同様のコードが生成されているものと思われる。
• (1)がCの標準の記述であり、(2)(3)(4)の方法で記述する理由はない。