图文精华

About云-梭伦科技»专题 › 交流区 › OpenStack系列 › Python问答 › 高性能的Python扩展（2）
分享

高性能的Python扩展（2）

xioaxu790 发表于 2014-11-27 20:58:57 [显示全部楼层] 回帖奖励 阅读模式 关闭右栏 0 14101
本帖最后由 pig2 于 2014-11-27 22:40 编辑
问题导读
1、你如何理解Wrold是存储N体状态的一个类？
2、如何利用PyArray_DATA宏来获得一个纸箱底层的内存？
3、在开始模拟时，N体被随机分配质量m，位置r和速度v的计算有哪些？






简介

这篇文章是这系列文章的第二篇，我们的关注点在使用Numpy API为Python编写C扩展模块的过程。在第一部分中，我们建立了一个简单的N体模拟，并发现其瓶颈是计算体之间的相互作用力，这是一个复杂度为O(N^2)的操作。通过在C语言中实现一个时间演化函数，我们大概能以大约70倍来加速计算。

如果你还没有看过第一篇文章，你应该在继续看这篇文章之前先看一下。


在这篇文章中，我们将牺牲我们代码中的一些通用性来提升性能。

回顾

Wrold是存储N体状态的一个类。我们的模拟将演化一系列时间步长下的状态。
  1. class World(object):

  2.     """World is a structure that holds the state of N bodies and

  3.     additional variables.

  4.  

  5.     threads : (int) The number of threads to use for multithreaded

  6.               implementations.

  7.     dt      : (float) The time-step.

  8.  

  9.     STATE OF THE WORLD: 

  10.  

  11.     N : (int) The number of bodies in the simulation.

  12.     m : (1D ndarray) The mass of each body.

  13.     r : (2D ndarray) The position of each body.

  14.     v : (2D ndarray) The velocity of each body.

  15.     F : (2D ndarray) The force on each body.

  16.  

  17.     TEMPORARY VARIABLES:

  18.  

  19.     Ft : (3D ndarray) A 2D force array for each thread's local storage.

  20.     s  : (2D ndarray) The vectors from one body to all others. 

  21.     s3 : (1D ndarray) The norm of each s vector. 

  22.  

  23.     NOTE: Ft is used by parallel algorithms for thread-local

  24.           storage. s and s3 are only used by the Python

  25.           implementation.

  26.     """

  27.     def __init__(self, N, threads=1, 

  28.                  m_min=1, m_max=30.0, r_max=50.0, v_max=4.0, dt=1e-3):

  29.         self.threads = threads

  30.         self.N  = N

  31.         self.m  = np.random.uniform(m_min, m_max, N)

  32.         self.r  = np.random.uniform(-r_max, r_max, (N, 2))

  33.         self.v  = np.random.uniform(-v_max, v_max, (N, 2))

  34.         self.F  = np.zeros_like(self.r)

  35.         self.Ft = np.zeros((threads, N, 2))

  36.         self.s  = np.zeros_like(self.r)

  37.         self.s3 = np.zeros_like(self.m)

  38.         self.dt = dt
复制代码


在开始模拟时，N体被随机分配质量m，位置r和速度v。对于每个时间步长，接下来的计算有：
1. 合力F，每个体上的合力根据所有其他体的计算。
2. 速度v，由于力的作用每个体的速度被改变。
3. 位置R，由于速度每个体的位置被改变。

访问宏

我们在第一部分创建的扩展模块使用宏来获取C语言中NumPy数组的元素。下面是这些宏中的一些宏的形式：

  1. #define r(x0, x1) (*(npy_float64*)((PyArray_DATA(py_r) +                \

  2.                                     (x0) * PyArray_STRIDES(py_r)[0] +   \

  3.                                     (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. static PyObject *evolve(PyObject *self, PyObject *args) {

  2.   // Declare variables. 

  3.   npy_int64 N, threads, steps, step, i, xi, yi;

  4.   npy_float64 dt;

  5.   PyArrayObject *py_m, *py_r, *py_v, *py_F;

  6.   npy_float64 *m, *r, *v, *F;

  7.  

  8.   // Parse arguments. 

  9.   if (!PyArg_ParseTuple(args, "ldllO!O!O!O!",

  10.                         &threads,

  11.                         &dt,

  12.                         &steps,

  13.                         &N,

  14.                         &PyArray_Type, &py_m,

  15.                         &PyArray_Type, &py_r,

  16.                         &PyArray_Type, &py_v,

  17.                         &PyArray_Type, &py_F)) {

  18.     return NULL;

  19.   }

  20.  

  21.   // Get underlying arrays from numpy arrays. 

  22.   m = (npy_float64*)PyArray_DATA(py_m);

  23.   r = (npy_float64*)PyArray_DATA(py_r);

  24.   v = (npy_float64*)PyArray_DATA(py_v);

  25.   F = (npy_float64*)PyArray_DATA(py_F);

  26.  

  27.   // Evolve the world. 

  28.   for(step = 0; step < steps; ++step) {

  29.     compute_F(N, m, r, F);

  30.  

  31.     for(i = 0; i < N; ++i) {

  32.       // Compute offsets into arrays. 

  33.       xi = 2 * i; 

  34.       yi = xi + 1;

  35.  

  36.       v[xi] += F[xi] * dt / m[i];

  37.       v[yi] += F[yi] * dt / m[i];

  38.  

  39.       r[xi] += v[xi] * dt;

  40.       r[yi] += v[yi] * dt;

  41.     }

  42.   }

  43.  

  44.   Py_RETURN_NONE;

  45. }
复制代码



从我们以前使用宏的实现来看，唯一的变化是我们必须明确地计算数组索引。我们可以将底层数组描述为二维`npy_float64`矩阵，但这需要一定的前期成本和内存管理。

计算力

与之前的实现一样，力以相同的方式进行计算。唯一的不同在符号，以及在for-loops循环中显式索引的计算。

  1. static inline void compute_F(npy_int64 N,

  2.                              npy_float64 *m,

  3.                              npy_float64 *r,

  4.                              npy_float64 *F) {

  5.   npy_int64 i, j, xi, yi, xj, yj;

  6.   npy_float64 sx, sy, Fx, Fy, s3, tmp;

  7.  

  8.   // Set all forces to zero. 

  9.   for(i = 0; i < N; ++i) {

  10.     F[2*i] = F[2*i + 1] = 0;

  11.   }

  12.  

  13.   // Compute forces between pairs of bodies.

  14.   for(i = 0; i < N; ++i) {

  15.     xi = 2*i;

  16.     yi = xi + 1;

  17.     for(j = i + 1; j < N; ++j) {

  18.       xj = 2*j; 

  19.       yj = xj + 1;

  20.  

  21.       sx = r[xj] - r[xi];

  22.       sy = r[yj] - r[yi];

  23.  

  24.       s3 = sqrt(sx*sx + sy*sy);

  25.       s3 *= s3*s3;

  26.  

  27.       tmp = m[i] * m[j] / s3;

  28.       Fx = tmp * sx;

  29.       Fy = tmp * sy;

  30.  

  31.       F[xi] += Fx;

  32.       F[yi] += Fy;

  33.       F[xj] -= Fx;

  34.       F[yj] -= Fy;

  35.     }

  36.   }

  37. }
复制代码



性能

我很惊讶地发现，相比于我们原来的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. static PyObject *evolve(PyObject *self, PyObject *args) {

  2.   // Variable declarations.

  3.   npy_int64 N, threads, steps, step, i;

  4.   npy_float64 dt;

  5.  

  6.   PyArrayObject *py_m, *py_r, *py_v, *py_F;

  7.   npy_float64 *m;

  8.   __m128d *r, *v, *F;

  9.  

  10.   // Parse arguments. 

  11.   if (!PyArg_ParseTuple(args, "ldllO!O!O!O!",

  12.                         &threads,

  13.                         &dt,

  14.                         &steps,

  15.                         &N,

  16.                         &PyArray_Type, &py_m,

  17.                         &PyArray_Type, &py_r,

  18.                         &PyArray_Type, &py_v,

  19.                         &PyArray_Type, &py_F)) {

  20.     return NULL;

  21.   }

  22.  

  23.   // Get underlying arrays from numpy arrays. 

  24.   m = (npy_float64*)PyArray_DATA(py_m);

  25.   r = (__m128d*)PyArray_DATA(py_r);

  26.   v = (__m128d*)PyArray_DATA(py_v);

  27.   F = (__m128d*)PyArray_DATA(py_F);

  28.  

  29.   // Evolve the world.

  30.   for(step = 0; step < steps; ++step) {

  31.     compute_F(N, m, r, F);

  32.  

  33.     for(i = 0; i < N; ++i) {

  34.       v[i] += F[i] * dt / m[i];

  35.       r[i] += v[i] * dt;

  36.     }

  37.   }

  38.  

  39.   Py_RETURN_NONE;

  40. }
复制代码



计算力

这里力的计算也比之前的实现更加紧凑。

  1. static inline void compute_F(npy_int64 N,

  2.                              npy_float64 *m,

  3.                              __m128d *r,

  4.                              __m128d *F) {

  5.   npy_int64 i, j;

  6.   __m128d s, s2, tmp;

  7.   npy_float64 s3;

  8.  

  9.   // Set all forces to zero.

  10.   for(i = 0; i < N; ++i) {

  11.     F[i] = _mm_set1_pd(0);

  12.   }

  13.  

  14.   // Compute forces between pairs of bodies. 

  15.   for(i = 0; i < N; ++i) {

  16.     for(j = i + 1; j < N; ++j) {

  17.       s = r[j] - r[i];

  18.  

  19.       s2 = s * s;

  20.       s3 = sqrt(s2[0] + s2[1]);

  21.       s3 *= s3 * s3;

  22.  

  23.       tmp = s * m[i] * m[j] / s3;

  24.       F[i] += tmp;

  25.       F[j] -= tmp;

  26.     }

  27.   }

  28. }
复制代码

性能
虽然这个明确的向量化代码更清晰，并比以前的版本更加紧凑&#8203;&#8203;，它的运行速度只提升了约1％，它每秒执行26427个时间步长。这可能是因为编译器能够使用相同的矢量指令优化之前的执行情况。

结论
如果全通用性是没有必要的，经常的性能提升，可以通过直接用C语言访问NumPy数组的底层内存来实现。
而在现在这个实例中，显式使用矢量内部没有提供多少好处，相对性能差异将在使用OpenMP的多芯设置中增大。


相关文章：高性能的Python扩展（1）



高性能的Python扩展（3）




关注公众号，获取大数据、人工智能20套、区块链资源5阶段等资源，随时更新，获取最新技术资源
回复

使用道具 举报

自定义广告语

没找到任何评论，期待你打破沉寂

高级模式
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册

本版积分规则

发表新帖
xioaxu790

金牌会员

关注

  • 851

    主题

  • 1000

    帖子

  • 108

    粉丝

TA的主题
关闭

推荐上一条 /2 下一条

关于我们· 联系我们· 加入我们· 小黑屋· 合作伙伴   

Copyright © 2001-2024 About云-梭伦科技 Powered by Discuz! X3.4 Licensed Discuz Team.

简书 / Smrz 京ICP备2020039040号 Wxb 简书网举报电话：021-34700000