簡介
這篇文章是這系列文章的第二篇，我們的關注點在使用Numpy API為Python編寫C擴充套件模組的過程。在第一部分中，我們建立了一個簡單的N體模擬，並發現其瓶頸是計算體之間的相互作用力，這是一個複雜度為O(N^2)的操作。通過在C語言中實現一個時間演化函式，我們大概能以大約70倍來加速計算。
如果你還沒有看過第一篇文章，你應該在繼續看這篇文章之前先看一下。
在這篇文章中，我們將犧牲我們程式碼中的一些通用性來提升效能。
回顧
Wrold是儲存N體狀態的一個類。我們的模擬將演化一系列時間步長下的狀態。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
class World(object): """World is a structure that holds the state of N bodies and additional variables. threads : (int) The number of threads to use for multithreaded implementations. dt : (float) The time-step. STATE OF THE WORLD: N : (int) The number of bodies in the simulation. m : (1D ndarray) The mass of each body. r : (2D ndarray) The position of each body. v : (2D ndarray) The velocity of each body. F : (2D ndarray) The force on each body. TEMPORARY VARIABLES: Ft : (3D ndarray) A 2D force array for each thread's local storage. s : (2D ndarray) The vectors from one body to all others. s3 : (1D ndarray) The norm of each s vector. NOTE: Ft is used by parallel algorithms for thread-local storage. s and s3 are only used by the Python implementation. """ def __init__(self, N, threads=1, m_min=1, m_max=30.0, r_max=50.0, v_max=4.0, dt=1e-3): self.threads = threads self.N = N self.m = np.random.uniform(m_min, m_max, N) self.r = np.random.uniform(-r_max, r_max, (N, 2)) self.v = np.random.uniform(-v_max, v_max, (N, 2)) self.F = np.zeros_like(self.r) self.Ft = np.zeros((threads, N, 2)) self.s = np.zeros_like(self.r) self.s3 = np.zeros_like(self.m) self.dt = dt |
在開始模擬時，N體被隨機分配質量m，位置r和速度v。對於每個時間步長，接下來的計算有：
1. 合力F，每個體上的合力根據所有其他體的計算。
2. 速度v，由於力的作用每個體的速度被改變。
3. 位置R，由於速度每個體的位置被改變。
訪問巨集
我們在第一部分建立的擴充套件模組使用巨集來獲取C語言中NumPy陣列的元素。下面是這些巨集中的一些巨集的形式：
|
1 2 3 |
#define r(x0, x1) (*(npy_float64*)((PyArray_DATA(py_r) + \ (x0) * PyArray_STRIDES(py_r)[0] + \ (x1) * PyArray_STRIDES(py_r)[1]))) |
像這樣使用巨集，我們能使用像r(i, j)這樣的簡單標記來訪問py_r陣列中的元素。不管陣列已經以某種形式被重塑或切片，索引值將匹配你在Python中看到的形式。
對於通用的程式碼，這就是我們想要的。在我們模擬的情況下，我們知道我們的陣列的特點：它們是連續的，並且從未被切片或重塑。我們可以利用這一點來簡化和提升我們程式碼的效能。
簡單的C擴充套件 2
在本節中，我們將看到一個修改版本的C擴充套件，它擯棄了訪問巨集和NumPy陣列底層資料的直接操作。本節中的程式碼src/simple2.c可在github上獲得。
為了進行比較，之前的實現也可在這裡獲得。
演化函式
從檔案的底部開始，我們可以看到，evolve函式與之前的版本一樣，以相同的方式解析Python引數，但現在我們利用PyArray_DATA巨集來獲得一個紙箱底層的記憶體。我們將這個指標命名為npy_float64，作為double的一個別名。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
static PyObject *evolve(PyObject *self, PyObject *args) { // Declare variables. npy_int64 N, threads, steps, step, i, xi, yi; npy_float64 dt; PyArrayObject *py_m, *py_r, *py_v, *py_F; npy_float64 *m, *r, *v, *F; // Parse arguments. if (!PyArg_ParseTuple(args, "ldllO!O!O!O!", &threads, &dt, &steps, &N, &PyArray_Type, &py_m, &PyArray_Type, &py_r, &PyArray_Type, &py_v, &PyArray_Type, &py_F)) { return NULL; } // Get underlying arrays from numpy arrays. m = (npy_float64*)PyArray_DATA(py_m); r = (npy_float64*)PyArray_DATA(py_r); v = (npy_float64*)PyArray_DATA(py_v); F = (npy_float64*)PyArray_DATA(py_F); // Evolve the world. for(step = 0; step < steps; ++step) { compute_F(N, m, r, F); for(i = 0; i < N; ++i) { // Compute offsets into arrays. xi = 2 * i; yi = xi + 1; v[xi] += F[xi] * dt / m[i]; v[yi] += F[yi] * dt / m[i]; r[xi] += v[xi] * dt; r[yi] += v[yi] * dt; } } Py_RETURN_NONE; } |
從我們以前使用巨集的實現來看，唯一的變化是我們必須明確地計算陣列索引。我們可以將底層陣列描述為二維npy_float64矩陣，但這需要一定的前期成本和記憶體管理。
計算力
與之前的實現一樣，力以相同的方式進行計算。唯一的不同在符號，以及在for-loops迴圈中顯式索引的計算。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
static inline void compute_F(npy_int64 N, npy_float64 *m, npy_float64 *r, npy_float64 *F) { npy_int64 i, j, xi, yi, xj, yj; npy_float64 sx, sy, Fx, Fy, s3, tmp; // Set all forces to zero. for(i = 0; i < N; ++i) { F[2*i] = F[2*i + 1] = 0; } // Compute forces between pairs of bodies. for(i = 0; i < N; ++i) { xi = 2*i; yi = xi + 1; for(j = i + 1; j < N; ++j) { xj = 2*j; yj = xj + 1; sx = r[xj] - r[xi]; sy = r[yj] - r[yi]; s3 = sqrt(sx*sx + sy*sy); s3 *= s3*s3; tmp = m[i] * m[j] / s3; Fx = tmp * sx; Fy = tmp * sy; F[xi] += Fx; F[yi] += Fy; F[xj] -= Fx; F[yj] -= Fy; } } } |
效能
我很驚訝地發現，相比於我們原來的C擴充套件，現在這種實現效能提升了45%。在相同的i5桌上型電腦上，以前的版本每秒執行17972個時間步長，而現在每秒執行26108個時間步長。這代表101倍提升了原始Python和NumPy的實現。
使用SIMD指令
在上面的實現中，我們在x和y方向上明確地計算出向量分量。如果我們願意放棄一些可移植性，並希望加快速度，我們可以使用x86的SIMD（單指令多資料）指令來簡化我們的程式碼。
本節中的程式碼src/simd1.c，可在github上獲得。
x86內部
在下面的程式碼中，我們將利用向量資料型別__m128d。數字128代表向量中的位元組數，而字母“d”表示向量分量的資料型別。在這種情況下，分量的型別是double（64位元組浮點）。其他向量的寬度和型別可根據不同的架構獲得。
多種內部函式都可以使用向量資料型別。英特爾在其網站上提供了一個方便的參考。
可移植性
我的工作環境是使用GNU gcc編譯器的Debian Wheezy系統。在這種環境下，定義可用的內部資料型別和函式的標頭檔案是x86intrin.h。這可能不能跨平臺移植。
演化函式
這裡的evolve函式比之前的版本更加緊湊。二維陣列轉換為__mm128d指標，並利用向量來排除顯式計算向量分量的需要。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
static PyObject *evolve(PyObject *self, PyObject *args) { // Variable declarations. npy_int64 N, threads, steps, step, i; npy_float64 dt; PyArrayObject *py_m, *py_r, *py_v, *py_F; npy_float64 *m; __m128d *r, *v, *F; // Parse arguments. if (!PyArg_ParseTuple(args, "ldllO!O!O!O!", &threads, &dt, &steps, &N, &PyArray_Type, &py_m, &PyArray_Type, &py_r, &PyArray_Type, &py_v, &PyArray_Type, &py_F)) { return NULL; } // Get underlying arrays from numpy arrays. m = (npy_float64*)PyArray_DATA(py_m); r = (__m128d*)PyArray_DATA(py_r); v = (__m128d*)PyArray_DATA(py_v); F = (__m128d*)PyArray_DATA(py_F); // Evolve the world. for(step = 0; step < steps; ++step) { compute_F(N, m, r, F); for(i = 0; i < N; ++i) { v[i] += F[i] * dt / m[i]; r[i] += v[i] * dt; } } Py_RETURN_NONE; } |
計算力
這裡力的計算也比之前的實現更加緊湊。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
static inline void compute_F(npy_int64 N, npy_float64 *m, __m128d *r, __m128d *F) { npy_int64 i, j; __m128d s, s2, tmp; npy_float64 s3; // Set all forces to zero. for(i = 0; i < N; ++i) { F[i] = _mm_set1_pd(0); } // Compute forces between pairs of bodies. for(i = 0; i < N; ++i) { for(j = i + 1; j < N; ++j) { s = r[j] - r[i]; s2 = s * s; s3 = sqrt(s2[0] + s2[1]); s3 *= s3 * s3; tmp = s * m[i] * m[j] / s3; F[i] += tmp; F[j] -= tmp; } } } |
效能
雖然這個明確的向量化程式碼更清晰，並比以前的版本更加緊湊，它的執行速度只提升了約1％，它每秒執行26427個時間步長。這可能是因為編譯器能夠使用相同的向量指令優化之前的執行情況。
結論
如果全通用性是沒有必要的，經常的效能提升，可以通過直接用C語言訪問NumPy陣列的底層記憶體來實現。
而在現在這個例項中，顯式使用向量內部沒有提供多少好處，相對效能差異將在使用OpenMP的多芯設定中增大。
下一部分
在本系列文章的下一部分，我們將利用OpenMP來並行使用這裡的實現，以及看如何利用多核心來擴充套件效能。
如果您有任何疑問，意見，建議或更正，請通過聯絡連結告訴我。