Question

利用我们刚刚学到的numpy，我们就可以轻松地把之前的纯Python代码矢量化。唯一需要引入的新功能是numpy的roll函数。该函数和我们手写的数组轮转技巧一样，但它运行在一个numpy数组上。一言以蔽之，它能将下面的数组轮转操作矢量化：

>>> import numpy as np
>>> np.roll([1,2,3,4], 1)
array([4, 1, 2, 3])

>>> np.roll([[1,2,3],[4,5,6]], 1, axis=1)
array([[3, 1, 2],
     [6, 4, 5]])

roll函数创建了一个新的numpy数组，这一点有利有弊。坏处是我们需要额外的时间开销来分配新的空间，然后还需要用合适的数据填满它。另一方面，一旦我们创建了这个轮转后的新数组，我们立刻就能够在其上进行矢量操作，而不会受到CPU缓存失效的影响。这可以极大影响我们对矩阵数据运算的速度。在本章后续，我们还将重写这部分代码来获得同样的好处而又不至于频繁分配更多内存。

有了这个额外的函数，我们就能将例6-6的Python扩散代码用更简单且矢量化的numpy数组重写。另外，我们将grid_xx和grid_yy的微分计算打散至另一个函数。例6-9展示了我们第一版的numpy扩散代码。

例6-9　初版numpy扩散

import numpy as np

grid_shape = (1024, 1024)

def laplacian(grid):
  return np.roll(grid, +1, 0) + np.roll(grid, -1, 0) + \
       np.roll(grid, +1, 1) + np.roll(grid, -1, 1) - 4 * grid

def evolve(grid, dt, D=1):
  return grid + dt * D * laplacian(grid)

def run_experiment(num_iterations):
  grid = np.zeros(grid_shape)

  block_low = int(grid_shape[0] * .4)
  block_high = int(grid_shape[0] * .5)
  grid[block_low:block_high, block_low:block_high] = 0.005

  start = time.time()
  for i in range(num_iterations):
    grid = evolve(grid, 0.1)
  return time.time() - start

我们会立刻发现这段代码短多了。这通常预示着高性能：我们将很多重活累活移到Python解释器之外，让模块内建的优化代码来帮我们解决特定的问题（但是，这必须要经过测试！）。我们的假设是numpy有更好的内存管理可以更快地给CPU提供所需的数据。但是，这一点实际取决于numpy的实现，让我们对代码进行性能分析来验证我们的假设。例6-10显示了结果。

例6-10　numpy 2阶扩散的性能指标

$ perf stat -e cycles,stalled-cycles-frontend,stalled-cycles-backend,instructions,\ cache-references,cache-misses,branches,branch-misses,task-clock,faults,\ minor-faults,cs,migrations -r 3 python diffusion_numpy.py Performance counter stats for 'python diffusion_numpy.py' (3 runs): 10,194,811,718 cycles # 3.332 GHz 4,435,850,419 stalled-cycles-frontend # 43.51% frontend cycles idle 2,055,861,567 stalled-cycles-backend # 20.17% backend cycles idle 15,165,151,844 instructions # 1.49 insns per cycle # 0.29 stalled cycles per insn 346,798,311 cache-references # 113.362 M/sec 519,793 cache-misses # 0.150 % of all cache refs 3,506,887,927 branches # 1146.334 M/sec 3,681,441 branch-misses # 0.10% of all branches 3059.219862 task-clock # 0.999 CPUs utilized 751,707 page-faults # 0.246 M/sec 751,707 minor-faults # 0.246 M/sec 8 context-switches # 0.003 K/sec 1 CPU-migrations # 0.000 K/sec 3.061883218 seconds time elapsed

这个结果显示出，相比于例6-8的降低内存分配的纯Python实现，仅仅改用numpy就给我们带来了40倍的速度提升。这是怎么做到的？

首先，我们要感谢numpy提供的矢量操作。虽然numpy版本看上去每个周期运行的指令较少，但是每条指令干了更多的活。也就是说，一条矢量操作指令可以对4个（或更多）数组元素进行乘法操作而不需要4条独立的乘法指令。最终导致我们解决同样问题的指令数变得更少。

numpy版本的函数能够降低指令总数还有其他原因。其中之一是纯Python版本需要完整的Python API可用，而numpy版本则没有这样的需求——比如说，纯Python矩阵只能在Python中附加，而不能在numpy模块内被附加。即使我们并没有显式使用这个功能（以及还有其他很多功能），仅仅是为了让系统确保它能够被使用也会带来额外的开销。由于numpy可以假定它存储的数据永远是数字，所有跟数组相关的操作都可以针对数字操作优化。等我们讨论Cython（7.6节）时，我们会为了性能继续移除各种必要功能，它甚至可以移除列表边界检查来加速整个列表查询。

通常，指令的数量跟性能并不相关——拥有较少指令的程序可能并没有高效地运行它们，或者它们可能都是慢速指令。不过，我们看到除了降低指令数量以外，numpy版本还减少了很多低效之处：缓存失效（仅0.15%，相比于之前的30.2%）。我们在6.3节中解释过，由于CPU必须等待从较慢的内存中读取数据而不是从缓存中直接使用，缓存失效降低了计算速度。事实上，内存碎片是如此决定性的性能因素，以至于在numpy中禁用矢量操作^[3]并保持其他一切不变的情况下，我们仍旧能看到一个相对纯Python版本极大的速度提升（例6-11）。

例6-11　禁用矢量操作的numpy 2阶扩散性能指标

$ perf stat -e cycles,stalled-cycles-frontend,stalled-cycles-backend,instructions,\ cache-references,cache-misses,branches,branch-misses,task-clock,faults,\ minor-faults,cs,migrations -r 3 python diffusion_numpy.py Performance counter stats for 'python diffusion_numpy.py' (3 runs): 48,923,515,604 cycles # 3.413 GHz 24,901,979,501 stalled-cycles-frontend # 50.90% frontend cycles idle 6,585,982,510 stalled-cycles-backend # 13.46% backend cycles idle 53,208,756,117 instructions # 1.09 insns per cycle # 0.47 stalled cycles per insn 83,436,665 cache-references # 5.821 M/sec 1,211,229 cache-misses # 1.452 % of all cache refs 4,428,225,111 branches # 308.926 M/sec 3,716,789 branch-misses # 0.08% of all branches 14334.244888 task-clock # 0.999 CPUs utilized 751,185 page-faults # 0.052 M/sec 751,185 minor-faults # 0.052 M/sec 24 context-switches # 0.002 K/sec 5 CPU-migrations # 0.000 K/sec 14.345794896 seconds time elapsed

这个结果展示出：当我们使用numpy的时候，给我们带来40倍速度提升的决定性因素并不是矢量操作指令集，而是内存的本地性以及减少的内存碎片。事实上，从之前的实验上我们可以看到，矢量操作仅给我们带来了40倍速度提升中的15%。^[4]

内存问题才是代码效率低下的决定性因素这一认知并不会令我们感到太过震惊。计算机就是专门被设计用来处理我们的问题的——把数字相乘和相加。瓶颈就取决于能否将这些数字以足够快的速度传输给CPU让它能够以最高的速度进行计算。

6.4.1　内存分配和就地操作

为了优化内存方面的影响，让我们试着使用跟例6-6同样的方法来降低我们numpy代码的内存分配次数。内存分配比我们之前讨论的缓存失效更糟。不仅仅要在缓存中找不到数据时去RAM中找到正确的数据，内存分配还必须向操作系统请求一块可用的数据并保留它。向操作系统进行请求所需的开销比简单填充缓存大很多——填充一次缓存失效是一个在主板上优化过的硬件行为，而内存分配则需要跟另一个进程、内核打交道来完成。

为了移除例6-9中的内存分配，我们会在代码开头预分配一些空间然后只使用就地操作。就地操作，比如+=、*=等，将其中一个输入重用为输出。这意味着我们不需要为计算的结果分配空间。

为了显式说明这点，我们看看当我们操作一个numpy数组时它的id如何改变（例6-12）。id是一个很好的追溯numpy数组的方式，因为id跟指向的内存区域有关。如果两个numpy数组具有相同的id，那么他们指向同一块内存区域。^[5]

例6-12　就地操作减少内存分配

>>> import numpy as np
>>> array1 = np.random.random((10,10))
>>> array2 = np.random.random((10,10))
>>> id(array1)
140199765947424 ❶
>>> array1 += array2
>>> id(array1)
140199765947424 # ❷
>>> array1 = array1 + array2
>>> id(array1)
140199765969792 # ❸

❶ ❷ 这两个id相同，因为我们使用了就地操作。这意味着array1的内存地址没有改变，我们只是改变了其内部的数据。

❸ 内存地址在这里改变了。在计算array1 + array2时，一个新的内存地址被分配并填入计算的结果。不过这样做也有好处，比如当原始数据需要被保留时（比如，array3 = array1 + array2允许你继续使用array1和array2，而就地操作销毁了部分原始数据）。

另外，我们看到一个由非就地操作带来的期望中的性能损失。对于小的numpy数组，这一开销大约占了整个计算时间的50%。至于大型计算，速度提升大概在几个百分点，但如果发生几百万次，这依然代表着一段很长的时间。在例6-13中，我们看到对小数组使用就地操作给我们带来了20%的速度提升。这一数字会随着数组长度增加而变得更多，因为内存分配会变得更繁重。

例6-13　使用就地操作减少内存分配

>>> %%timeit array1, array2 = np.random.random((10,10)), np.random.random((10,10))
# ❶
... array1 = array1 + array2
...
100000 loops, best of 3: 3.03 us per loop

>>> %%timeit array1, array2 = np.random.random((10,10)), np.random.random((10,10))
... array1 += array2
...
100000 loops, best of 3: 2.42 us per loop

❶ 注意我们使用了%%timeit而不是%timeit，这让我们可以指定用于设置环境的代码免于时间测量。

虽然改写我们例6-9中的代码并不十分复杂，但就地操作的坏处在于它确实会让代码有一点难以阅读。在例6-14中我们可以看到这一重构的结果。我们初始化grid和next_grid矩阵，然后重复交换两者。grid是当前系统的已知信息，在运行了evolve之后，next_grid包含了更新后的信息。

例6-14　将大多数numpy操作改为就地操作

def laplacian(grid, out):
  np.copyto(out, grid)
  out *= -4
  out += np.roll(grid, +1, 0)
  out += np.roll(grid, -1, 0)
  out += np.roll(grid, +1, 1)
  out += np.roll(grid, -1, 1)

def evolve(grid, dt, out, D=1):
  laplacian(grid, out)
  out *= D * dt
  out += grid

def run_experiment(num_iterations):
  next_grid = np.zeros(grid_shape)
  grid = np.zeros(grid_shape)

  block_low = int(grid_shape[0] * .4)
  block_high = int(grid_shape[0] * .5)
  grid[block_low:block_high, block_low:block_high] = 0.005

  start = time.time()
  for i in range(num_iterations):
    evolve(grid, 0.1, next_grid)
    grid, next_grid = next_grid, grid # ❶
  return time.time() - start

❶ 因为evolve的输出被保存在输出矩阵next_grid，我们必须为下一次迭代交换这两个变量来让grid包含最新的信息。这一交换操作非常便宜，因为仅仅互换了数据的引用，而不是数据本身。

记住因为我们想要让每个操作都就地完成，所以每一个矢量操作都必须在同一行完成。这可以让一句简单的A = A * B + C变得相当费解。因为Python非常强调可读性，所以我们必须确保这些改动给我们带来的速度提升足以弥补这一点。

比较例6-15和例6-10的性能指标，我们看到移除不需要的内存分配给我们的代码带来了29%的速度提升。部分是因为降低了缓存失效，但主要还是因为降低了内存缺页。

例6-15　numpy就地操作内存的性能指标

$ perf stat -e cycles,stalled-cycles-frontend,stalled-cycles-backend,instructions,\ cache-references,cache-misses,branches,branch-misses,task-clock,faults,\ minor-faults,cs,migrations -r 3 python diffusion_numpy_memory.py Performance counter stats for 'python diffusion_numpy_memory.py' (3 runs): 7,864,072,570 cycles # 3.330 GHz 3,055,151,931 stalled-cycles-frontend # 38.85% frontend cycles idle 1,368,235,506 stalled-cycles-backend # 17.40% backend cycles idle 13,257,488,848 instructions # 1.69 insns per cycle # 0.23 stalled cycles per insn 239,195,407 cache-references # 101.291 M/sec 2,886,525 cache-misses # 1.207 % of all cache refs 3,166,506,861 branches # 1340.903 M/sec 3,204,960 branch-misses # 0.10% of all branches 2361.473922 task-clock # 0.999 CPUs utilized 6,527 page-faults # 0.003 M/sec 6,527 minor-faults # 0.003 M/sec 6 context-switches # 0.003 K/sec 2 CPU-migrations # 0.001 K/sec 2.363727876 seconds time elapsed

6.4.2　选择优化点：找到需要被修正的地方

从例6-14的代码中，似乎可以看到我们解决了手头大部分的问题：我们使用numpy降低了CPU负担，我们降低了解决问题所必要的内存分配次数。但是，永远有更多的调查等待着我们。如果我们对代码做一个逐行性能分析（例6-16），就会看到我们大部分的工作都在laplacian函数中完成。事实上evolve函数有93%的时间花费在laplacian函数中。

例6-16　逐行性能分析显示laplacian花费了太多时间

Wrote profile results to diffusion_numpy_memory.py.lprof
Timer unit: 1e-06 s

File: diffusion_numpy_memory.py
Function: laplacian at line 8
Total time: 3.67347 s

Line #   Hits    Time   Per Hit  % Time  Line Contents
==============================================================
   8                     @profile
   9                     def laplacian(grid, out):
  10    500  162009   324.0   4.4    np.copyto(out, grid)
  11    500  111044   222.1   3.0    out *= -4
  12    500  464810   929.6  12.7    out += np.roll(grid, +1, 0)
  13    500  432518   865.0  11.8    out += np.roll(grid, -1, 0)
  14    500   1261692  2523.4  34.3    out += np.roll(grid, +1, 1)
  15    500   124139   2482.8  33.8    out += np.roll(grid, -1, 1)

File: diffusion_numpy_memory.py
Function: evolve at line 17
Total time: 3.97768 s

Line #   Hits    Time   Per Hit  % Time  Line Contents
==============================================================
  17                     @profile
  18                     def evolve(grid, dt, out, D=1):
  19    500  3691674   7383.3  92.8   laplacian(grid, out)
  20    500   111687  223.4   2.8   out *= D * dt
  21    500   174320  348.6   4.4   out += grid

laplacian函数如此慢的原因可能有很多。但是，这里我们主要考虑两个高级因素。首先，在调用np.roll时会为创建新的矢量分配内存（我们可以通过查看该函数的说明文档来验证这一点）。这意味着即使我们在之前的重构中移除了7处内存分配，仍然有4处分配是我们遗漏的。另外，np.roll是一个非常通用的函数，有很多代码用来处理特殊情况。既然我们已经完全清楚我们需要做什么（那就是将每个维度的第一列的数据移到最后一列），那么我们就可以重写这个函数来消除大部分无用的代码。我们甚至可以将np.roll的代码逻辑和add操作合并，生成一个高度专用的roll_add函数以最少的内存分配次数和最精简的逻辑来完成我们的任务。

例6-17显示了这个重构的细节。我们需要的仅是创建新的roll_add函数并让laplacian使用它。因为numpy支持复杂索引（fancy indexing），实现这样一个函数只需要保证不搞乱索引就行。但是，之前说过，这样的代码性能虽好，可读性却不怎么样。

　警告　

注意我们除了彻底测试之外，还在docstring上花费了额外的维护工作。当你在做类似这样的优化时，维护代码的可读性是非常重要的，这些步骤是为了确保你的代码以你期望的方式工作，让未来的程序开发者能够维护你的代码并在出问题时清楚地知道你的代码做了些什么。

例6-17　创建我们自己的轮转函数

import numpy as np

def roll_add(rollee, shift, axis, out):
  """
  Given a matrix, rollee, and an output matrix, out, this function will
  perform the calculation:

    >>> out += np.roll(rollee, shift, axis=axis)

  This is done with the following assumptions:
    * rollee is 2D
    * shift will only ever be +1 or -1
    * axis will only ever be 0 or 1 (also implied by the first assumption)

Using these assumptions, we are able to speed up this function by avoiding
extra machinery that numpy uses to generalize the `roll` function and also
by making this operation intrinsically in-place.
"""
if shift == 1 and axis == 0:
  out[1:, :] += rollee[:-1, :]
  out[0 , :] += rollee[-1, :]
elif shift == -1 and axis == 0:
  out[:-1, :] += rollee[1:, :]
  out[-1 , :] += rollee[0, :]
elif shift == 1 and axis == 1:
  out[:, 1:] += rollee[:, :-1]
  out[:, 0 ] += rollee[:, -1]
elif shift == -1 and axis == 1:
  out[:, :-1] += rollee[:, 1:]
  out[:, -1] += rollee[:, 0]

def test_roll_add():
  rollee = np.asarray([[1,2],[3,4]])
  for shift in (-1, +1):
    for axis in (0, 1):
      out = np.asarray([[6,3],[9,2]])
      expected_result = np.roll(rollee, shift, axis=axis) + out
      roll_add(rollee, shift, axis, out)
      assert np.all(expected_result == out)

def laplacian(grid, out):
  np.copyto(out, grid)
  out *= -4
  roll_add(grid, +1, 0, out)
  roll_add(grid, -1, 0, out)
  roll_add(grid, +1, 1, out)
  roll_add(grid, -1, 1, out)

如果我们查看例6-18中本次重写后的性能指标，我们会发现虽然它比例6-14有一个很大的性能提升（实际上提升了70%），但是大多数性能指标都是相同的。Page-faults的数量下降了，但是下降了不到70%。同样，cache-misses跟cache-references也下降了，但是也不足以导致整体这么大的性能提升。这里最重要的参数是instructions指标。Instructions指标记录了CPU需要执行多少指令来完成整个程序——也就是说，CPU需要干多少事情。使用特化的roll_add函数后instructions指标降低了2.86倍。这是因为我们不再需要依靠numpy提供的完整功能来轮转我们的矩阵，而是利用我们对数据的了解（我们的数据只有两个维度且只需轮转一次）创建了更精简的功能。我们将在7.6节中继续讨论关于在numpy和Python中精简不需要功能的主题。

例6-18　numpy使用就地内存操作和特化laplacian函数的性能指标

$ perf stat -e cycles,stalled-cycles-frontend,stalled-cycles-backend,instructions,\
  cache-references,cache-misses,branches,branch-misses,task-clock,faults,\
  minor-faults,cs,migrations -r 3 python diffusion_numpy_memory2.py

Performance counter stats for 'python diffusion_numpy_memory2.py' (3 runs):

   4,303,799,244 cycles           #   3.108 GHz
   2,814,678,053 stalled-cycles-frontend  #  65.40% frontend cycles idle
   1,635,172,736 stalled-cycles-backend   #  37.99% backend cycles idle
   4,631,882,411 instructions       #  1.08 insns per cycle
                      #  0.61 stalled cycles per insn
  272,151,957 cache-references      # 196.563 M/sec
    2,835,948 cache-misses        # 1.042 % of all cache refs
  621,565,054 branches          # 448.928 M/sec
    2,905,879 branch-misses       #   0.47% of all branches
  1384.555494 task-clock        #   0.999 CPUs utilized
      5,559 page-faults         #   0.004 M/sec
      5,559 minor-faults        #   0.004 M/sec
        6 context-switches      #   0.004 K/sec
        3 CPU-migrations      #   0.002 K/sec
  1.386148918 seconds time elapsed