高性能的Python扩展（3）

问题导读
1、如何重新实现compute_F函数？
2、性能测试，需要有哪些环境？
3、计算模拟中每对体之间的相互作用力，其复杂度为多少？






简介
本文是这个系列的第三篇，我们关注于使用NumPy API为Python编写高性能的C扩展模块。在本文中，我们将使用OpenMP来并行第二部分中的实现。

回顾
Wrold是存储N体状态的一个类。我们的模拟将演化一系列时间步长下的状态。

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，由于速度每个体的位置被改变。

计算力：串行代码

下面是之前文章实现中（全部的源代码在这里）的compute_F函数。这个函数计算模拟中每对体之间的相互作用力，其复杂度为O(N^2)。

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;
    }
  }
}
复制代码

性能
使用这一系列的实现，我们的代码在i5台式机上能每秒执行26427个时间步长。

计算力：并行代码
下面是重新实现的compute_F函数，它使用OpenMP来并行循环。完整的源文件src/simdomp1.c可以在github上获得。

static inline void compute_F(npy_int64 threads,                            
                             npy_int64 N,
                             npy_float64 *m,
                             __m128d *r,
                             __m128d *F, 
                             __m128d *Ft) {
  npy_int64 id, i, j, Nid;
  __m128d s, s2, tmp;
  npy_float64 s3;
 
#pragma omp parallel private(id, i, j, s, s2, s3, tmp, Nid) 
  {
    id = omp_get_thread_num();
    Nid = N * id; // Zero-index in thread-local array Ft. 
 
    // Zero out the thread-local force arrays. 
    for(i = 0; i < N; i++) {
      Ft[Nid + i] = _mm_set1_pd(0);
    }
 
    // Compute forces between pairs of bodies.
#pragma omp for schedule(dynamic, 8) 
    for(i = 0; i < N; ++i) {
      F[i] = _mm_set1_pd(0);
 
      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;
 
        Ft[Nid + i] += tmp;
        Ft[Nid + j] -= tmp;
      }
    }
 
    // Sum the thread-local forces computed above.
#pragma omp for
    for(i = 0; i < N; ++i) {
      for(id = 0; id < threads; ++id) {
        F[i] += Ft[N*id + i];
      }
    }
  }
}
复制代码

线程局部存储

串行版本和并行版本之间最明显的不同在于Ft数组的出现。Ft数组通过每个同时运行的线程为力的计算提供线程局部存储。这些局部的数组然后在一个单独的循环中相加，从而得到一个力F的数组。 局部存储对于避免竞争条件是有必要的。可以想象一下，如果一个线程正在执行i=0和j=1，而此时另一个线程正在使用i=1和j=2。在我们的算法中，它们都试图修改F[1]，这导致了一种竞争。 你可以在维基百科上阅读有关竞争条件（race conditions）的内容。

OpenMP

OpenMP API通过使用如omp_get_thread_num这样的函数和如#pragma omp parallel这样的指令来执行。 在上面的代码中，我们创建了许多有#pragma omp parallel指令的线程。private指令列出了每个线程的私有变量。所有其他变量是公有的。 #pragma omp指令允许for循环并行进行，每个线程处理不同的指数。我们使用schedule(dynamic, 8)指令来告诉编译器使用8块大小的动态调度。当每个循环花费的时间不同时，动态调度是有用的，正如这里的这种情况。 需要注意的是，在每个并行for循环的末尾有一个隐式屏障。所有的线程在任何可以提前之前都必须完成循环。 在OpenMP 网站上参见有关该API的更多信息。

性能

因为我们增加了一个额外的for循环来相加局部线程的力的数组，我们可以预期单核性能会稍受损失。在我们的测试中，在使用单线程时，OpenMP版本的运行速度比没有OpenMP的版本慢约3%。 在Intel i5台式机上，使用四核及四线程，OpenMP使用SIMD指令执行，每秒80342个时间步长。这比我们原来使用NumPy的Python实现块312倍，比我们最快的串行实现快3倍。

扩展性


上图显示OpenMP代码是如何在四核i5台式机上扩展线程数量的。标有”OMP+SIMD”的点对应的是本文中的实现。一个类似的版本可以在这里获得，没有明确的向量指令标有”OMP”。使用Python和NumPy的原始版本被用于比较（左下）。 性能可以很好的达到物理核的数量，并大幅减少额外线程的增加。请注意，这不是一般的事实，所以使用不同数量的线程来衡量你的代码总是一个好主意。

环境变量
这里有许多环境变量可以改变使用OpenMp程序的行为。这些都对性能有重大的影响，特别是对有多个物理处理器的计算机。 请参阅GNU libgomp 文档了解详细信息。

性能总结：所有实现

这里是这个系列文章中所有实现版本的基准。这个运行在一个四核Intel i5台式机上，在Debian Wheezy系统下使用gcc 4.7.2。

实现	时间步长/每秒	加速
Python+Numpy	257	1
Simple C 1	17974	70
Simple C 2	26136	102
SIMD	26423	103
OMP 4 cores	76657	298
OMP+SIMD 4 cores	80342	313

附加测试

DigitalOcean20核虚拟机
在DigitalOcean20核虚拟机实例上运行上面相同的最快的实现。3个因素的最好的性能比在四核四线程的i5台式机上运行时的性能差。我不知道这个差异的原因是否是因为虚拟化和共享内核。


RunAbove176线程虚拟机
RunAbove对他们IBM Power 8架构免费一个月的使用进行推广。为了好玩，我在他们176线程的实例上运行了相同的实现。之前我从来没有使用过PowerPC架构，我很高兴代码运行不需要任何修改。




Cython

Matthew Honnibal提交了一个Cythond的实现，可以在这里获得。这个实现和我们当初在第一部分中的C语言实现有些类似。


相关文章：

高性能的Python扩展（1）

高性能的Python扩展（2）