高性能的Python扩展(3),高性能python扩展,未经许可,禁止转载!英文


本文由 编橙之家 - Janzou 翻译,黄利民 校稿。未经许可,禁止转载!
英文出处:crumpington。欢迎加入翻译组。
  • 高性能的Python扩展(1)
  • 高性能的Python扩展(2)

简介

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

回顾

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

Python
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)。

Python
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上获得。

Python
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

Python

			
				
			
		

附加测试

DigitalOcean20核虚拟机

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

RunAbove176线程虚拟机

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

Cython

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

结论

当性能至关重要时,一个C语言扩展可以显著提升性能。使用SIMD指令可以使性能受益并简化代码,但有移植的成本。OpenMP的支持可以轻松地增加现有代码,并能提高在多核系统中的性能。 如果您有任何疑问,意见,建议或更正,请通过联系链接告诉我。

评论关闭