为何numpy.sum比经优化(自动向量化)的C循环更快?
我尝试编写一个在double数组上性能与numpy.sum相当的C程序,但未能成功。以下是我的测试过程:
NumPy性能测试代码
import numpy as np import time SIZE=4000000 REPS=5 xs = np.random.rand(SIZE) print(xs.dtype) for _ in range(REPS): start = time.perf_counter() r = np.sum(xs) end = time.perf_counter() print(f"{SIZE / (end-start) / 10**6:.2f} MFLOPS ({r:.2f})")
测试输出
float64 2941.61 MFLOPS (2000279.78) 3083.56 MFLOPS (2000279.78) 3406.18 MFLOPS (2000279.78) 3712.33 MFLOPS (2000279.78) 3661.15 MFLOPS (2000279.78)
C语言测试代码
#include <stdio.h> #include <stdlib.h> #include <time.h> #define SIZE 4000000 #define REPS 5 double *make_random_array(long array_size) { double *array = malloc(array_size * sizeof(double)); if (array == NULL) return NULL; srand(0); for (size_t i = 0; i < array_size; ++i) { array[i] = (double)rand() / RAND_MAX; } return array; } double sum_array(const double *array, long size) { double sum = 0.0; for (size_t i = 0; i < size; ++i) { sum += array[i]; } return sum; } int main() { double *xs = make_random_array(SIZE); if (xs == NULL) return 1; for (int i = 0; i < REPS; i++) { clock_t start_time = clock(); double r = sum_array(xs, SIZE); clock_t end_time = clock(); double dt = (double)(end_time - start_time) / CLOCKS_PER_SEC; printf("%.2f MFLOPS (%.2f)\n", (double)SIZE / dt / 1000000, r); } free(xs); return 0; }
编译命令及输出
使用gcc -o main -Wall -O3 -mavx main.c编译后,输出如下:
1850.14 MFLOPS (1999882.86) 1857.01 MFLOPS (1999882.86) 1900.24 MFLOPS (1999882.86) 1903.86 MFLOPS (1999882.86) 1906.58 MFLOPS (1999882.86)
额外观察
- C代码性能明显低于NumPy
- 通过
top观察到Python进程CPU使用率约100%,说明NumPy未进行并行计算 - C代码已使用256位AVX寄存器(编译为汇编可见
vaddsd指令),本机CPU为13th Gen Intel(R) Core(TM) i9-13900K,不支持AVX-512:
$ egrep 'model name|flags' /proc/cpuinfo | head -n2 model name : 13th Gen Intel(R) Core(TM) i9-13900K flags : fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb intel_pt sha_ni xsaveopt xsavec xgetbv1 xsaves split_lock_detect avx_vnni dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req hfi umip pku ospke waitpkg gfni vaes vpclmulqdq tme rdpid movdiri movdir64b fsrm md_clear serialize pconfig arch_lbr ibt flush_l1d arch_capabilities
请问NumPy采用了什么技术技巧,使得其性能超过这段C代码?
NumPy能超过你编写的C代码,核心是以下几个技术点:
1. 高效的SIMD向量化优化
你的C代码虽然用了-mavx编译,但生成的是vaddsd指令——这是单double值加法,每次仅处理1个元素。而NumPy内部会使用vaddpd这类打包double加法指令,利用AVX2的256位寄存器一次处理4个double值(每个double占8字节,256/8=4),理论吞吐量是单指令的4倍。
gcc默认对简单逐元素累加循环的向量化会比较保守,可能担心浮点累加顺序影响结果(尽管sum的顺序不影响IEEE合规性,但编译器会做保守处理)。而NumPy的sum实现是专门针对向量化编写的,明确使用打包指令批量处理数据。
2. 循环展开与软件流水线优化
NumPy会对循环进行多轮展开(比如一次处理8/16个元素),结合软件流水线技术,让CPU的指令执行单元尽可能被填满,减少流水线停顿。你的C代码是简单单循环,gcc的O3优化可能会做部分展开,但不如NumPy手动优化的彻底。
3. 内存对齐优化
NumPy的数组默认按CPU缓存行对齐(通常64字节),而你用malloc分配的数组无法保证对齐。不对齐的内存访问会导致额外CPU开销,尤其是SIMD操作时需要拆分内存访问,降低效率。你可以改用posix_memalign分配对齐内存,测试性能变化。
4. 针对性CPU架构优化
NumPy编译时通常会启用-march=native或针对目标CPU的特定优化,充分利用你i9-13900K支持的AVX2、FMA等指令集特性。而你编译C代码时仅加了-mavx,未启用AVX2(从CPU flags可看到支持),限制了gcc生成更高效的指令。尝试用gcc -o main -Wall -O3 -mavx2 -mfma main.c编译,性能会有明显提升。
5. 分段累加策略优化
NumPy的sum实现会用分段累加:先将数组分成多个块,每个块用独立寄存器累加,最后合并所有块的结果。这种方式既提升了SIMD利用率,又减少了单寄存器累加的依赖链停顿(CPU需等待前一次加法完成才能执行下一次,分段累加可并行处理多个块)。而你的代码用单个sum变量逐元素累加,依赖链较长,限制了CPU乱序执行能力。
内容的提问来源于stack exchange,提问作者dnalor

