为什么 numba popcount 代码的速度是等效 C 代码的两倍?

Rap*_*ael 7 c python performance x86-64 numba

我有这个简单的 python/numba 代码:

\n
from numba import njit\nimport numba as nb\n@nb.njit(nb.uint64(nb.uint64))\ndef popcount(x): \n      b=0\n      while(x > 0):\n          x &= x - nb.uint64(1)   \n          b+=1\n      return b\n@njit\ndef timed_loop(n):\n    summand = 0\n    for i in range(n):\n        summand += popcount(i)\n    return summand\n
Run Code Online (Sandbox Code Playgroud)\n

它只是将整数 0 到 n - 1 的 popcount 相加。

\n

当我计时时,我得到:

\n
%timeit timed_loop(1000000)\n340 \xc2\xb5s \xc2\xb1 1.08 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 1,000 loops each)\n
Run Code Online (Sandbox Code Playgroud)\n

事实证明,llvm 巧妙地将 popcount 函数转换为本机 CPU POPCNT 指令,因此我们应该期望它很快。但问题是,速度有多快。

\n

我想将它与 C 版本进行比较,看看速度差异。

\n
 #include <stdio.h>\n#include <time.h>\n\n// Function to calculate the population count (number of set bits) of an integer using __builtin_popcount\nint popcount(int num) {\n    return __builtin_popcount(num);\n}\n\nint main() {\n    unsigned int n;\n    printf("Enter the value of n: ");\n    scanf("%d", &n);\n\n    // Variables to store start and end times\n    struct timespec start_time, end_time;\n\n    // Get the current time as the start time\n    clock_gettime(CLOCK_MONOTONIC, &start_time);\n\n    int sum = 0;\n    for (unsigned int i = 0; i < n; i++) {\n        sum += popcount(i);\n    }\n\n    // Get the current time as the end time\n    clock_gettime(CLOCK_MONOTONIC, &end_time);\n\n    // Calculate the elapsed time in microseconds\n    long long elapsed_time = (end_time.tv_sec - start_time.tv_sec) * 1000000LL +\n                            (end_time.tv_nsec - start_time.tv_nsec) / 1000;\n\n    printf("Sum of population counts from 0 to %d-1 is: %d\\n", n, sum);\n    printf("Elapsed time: %lld microseconds\\n", elapsed_time);\n\n    return 0;\n}\n
Run Code Online (Sandbox Code Playgroud)\n

然后我用 编译了这个-march=native -Ofast。我尝试了 gcc 和 clang,结果非常相似。

\n
./popcount \nEnter the value of n: 1000000\nSum of population counts from 0 to 1000000-1 is: 9884992\nElapsed time: 732 microseconds\n
Run Code Online (Sandbox Code Playgroud)\n

为什么 numba 的速度是 C 代码的两倍?

\n

Jér*_*ard 9

TL;DR: GCC 和 Clang 版本之间的性能差距是由于使用标量指令与 SIMD 指令造成的。Numba 和 Clang 版本之间的性能差距来自于两个版本之间的整数大小不同:64 位与 32 位


绩效结果

首先,我也能够在我的 Intel i5-9600KF 上重现该问题。以下是结果(和版本):

Numba 0.56.4:  170.089 ms
Clang 14.0.6:  190.350 ms
GCC 12.2.0:    328.133 ms
Run Code Online (Sandbox Code Playgroud)

为了了解发生了什么,我们需要分析所有编译器生成的汇编代码。


汇编代码

以下是 GCC 生成的热循环的汇编代码:

.L5:
    xorl    %edx, %edx
    popcntl %eax, %edx
    incl    %eax
    addl    %edx, %ebx
    cmpl    %ecx, %eax
    jne .L5
Run Code Online (Sandbox Code Playgroud)

这是 Clang 制作的:

.LBB1_3:                                # =>This Inner Loop Header: Depth=1
    vpand   %ymm5, %ymm0, %ymm12
    vpshufb %ymm12, %ymm6, %ymm12
    vpsrlw  $4, %ymm0, %ymm13
    vpand   %ymm5, %ymm13, %ymm13
    vpshufb %ymm13, %ymm6, %ymm13
    vpaddb  %ymm12, %ymm13, %ymm12
    vpunpckhdq  %ymm1, %ymm12, %ymm13   # ymm13 = ymm12[2],ymm1[2],ymm12[3],ymm1[3],ymm12[6],ymm1[6],ymm12[7],ymm1[7]
    vpsadbw %ymm1, %ymm13, %ymm13
    vpunpckldq  %ymm1, %ymm12, %ymm12   # ymm12 = ymm12[0],ymm1[0],ymm12[1],ymm1[1],ymm12[4],ymm1[4],ymm12[5],ymm1[5]
    vpsadbw %ymm1, %ymm12, %ymm12
    vpackuswb   %ymm13, %ymm12, %ymm12
    vpaddd  %ymm2, %ymm0, %ymm13
    vpaddd  %ymm12, %ymm8, %ymm8
    vpand   %ymm5, %ymm13, %ymm12
    vpshufb %ymm12, %ymm6, %ymm12
    vpsrlw  $4, %ymm13, %ymm13
    vpand   %ymm5, %ymm13, %ymm13
    vpshufb %ymm13, %ymm6, %ymm13
    vpaddb  %ymm12, %ymm13, %ymm12
    vpunpckhdq  %ymm1, %ymm12, %ymm13   # ymm13 = ymm12[2],ymm1[2],ymm12[3],ymm1[3],ymm12[6],ymm1[6],ymm12[7],ymm1[7]
    vpsadbw %ymm1, %ymm13, %ymm13
    vpunpckldq  %ymm1, %ymm12, %ymm12   # ymm12 = ymm12[0],ymm1[0],ymm12[1],ymm1[1],ymm12[4],ymm1[4],ymm12[5],ymm1[5]
    vpsadbw %ymm1, %ymm12, %ymm12
    vpackuswb   %ymm13, %ymm12, %ymm12
    vpaddd  %ymm3, %ymm0, %ymm13
    vpaddd  %ymm12, %ymm9, %ymm9
    vpand   %ymm5, %ymm13, %ymm12
    vpshufb %ymm12, %ymm6, %ymm12
    vpsrlw  $4, %ymm13, %ymm13
    vpand   %ymm5, %ymm13, %ymm13
    vpshufb %ymm13, %ymm6, %ymm13
    vpaddb  %ymm12, %ymm13, %ymm12
    vpunpckhdq  %ymm1, %ymm12, %ymm13   # ymm13 = ymm12[2],ymm1[2],ymm12[3],ymm1[3],ymm12[6],ymm1[6],ymm12[7],ymm1[7]
    vpsadbw %ymm1, %ymm13, %ymm13
    vpunpckldq  %ymm1, %ymm12, %ymm12   # ymm12 = ymm12[0],ymm1[0],ymm12[1],ymm1[1],ymm12[4],ymm1[4],ymm12[5],ymm1[5]
    vpsadbw %ymm1, %ymm12, %ymm12
    vpackuswb   %ymm13, %ymm12, %ymm12
    vpaddd  %ymm4, %ymm0, %ymm13
    vpaddd  %ymm12, %ymm10, %ymm10
    vpand   %ymm5, %ymm13, %ymm12
    vpshufb %ymm12, %ymm6, %ymm12
    vpsrlw  $4, %ymm13, %ymm13
    vpand   %ymm5, %ymm13, %ymm13
    vpshufb %ymm13, %ymm6, %ymm13
    vpaddb  %ymm12, %ymm13, %ymm12
    vpunpckhdq  %ymm1, %ymm12, %ymm13   # ymm13 = ymm12[2],ymm1[2],ymm12[3],ymm1[3],ymm12[6],ymm1[6],ymm12[7],ymm1[7]
    vpsadbw %ymm1, %ymm13, %ymm13
    vpunpckldq  %ymm1, %ymm12, %ymm12   # ymm12 = ymm12[0],ymm1[0],ymm12[1],ymm1[1],ymm12[4],ymm1[4],ymm12[5],ymm1[5]
    vpsadbw %ymm1, %ymm12, %ymm12
    vpackuswb   %ymm13, %ymm12, %ymm12
    vpaddd  %ymm12, %ymm11, %ymm11
    vpaddd  %ymm7, %ymm0, %ymm0
    addl    $-32, %edx
    jne .LBB1_3
Run Code Online (Sandbox Code Playgroud)

这是 Numba 制作的:

.LBB0_8:
    vpand   %ymm0, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpsrlw  $4, %ymm0, %ymm7
    vpand   %ymm7, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpaddb  %ymm6, %ymm7, %ymm6
    vpaddq  -40(%rsp), %ymm0, %ymm7
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm1, %ymm1
    vpand   %ymm7, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpsrlw  $4, %ymm7, %ymm7
    vpand   %ymm7, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpaddb  %ymm6, %ymm7, %ymm6
    vpaddq  -72(%rsp), %ymm0, %ymm7
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm2, %ymm2
    vpand   %ymm7, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpsrlw  $4, %ymm7, %ymm7
    vpand   %ymm7, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpaddb  %ymm6, %ymm7, %ymm6
    vpaddq  %ymm0, %ymm8, %ymm7
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm3, %ymm3
    vpand   %ymm7, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpsrlw  $4, %ymm7, %ymm7
    vpand   %ymm7, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpaddb  %ymm6, %ymm7, %ymm6
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm4, %ymm4
    vpaddq  %ymm0, %ymm11, %ymm6
    vpand   %ymm6, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpsrlw  $4, %ymm6, %ymm6
    vpand   %ymm6, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpaddb  %ymm7, %ymm6, %ymm6
    vpaddq  %ymm0, %ymm12, %ymm7
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm1, %ymm1
    vpand   %ymm7, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpsrlw  $4, %ymm7, %ymm7
    vpand   %ymm7, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpaddb  %ymm6, %ymm7, %ymm6
    vpaddq  %ymm0, %ymm13, %ymm7
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm2, %ymm2
    vpand   %ymm7, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpsrlw  $4, %ymm7, %ymm7
    vpand   %ymm7, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpaddb  %ymm6, %ymm7, %ymm6
    vpaddq  %ymm0, %ymm14, %ymm7
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm3, %ymm3
    vpand   %ymm7, %ymm9, %ymm6
    vpshufb %ymm6, %ymm10, %ymm6
    vpsrlw  $4, %ymm7, %ymm7
    vpand   %ymm7, %ymm9, %ymm7
    vpshufb %ymm7, %ymm10, %ymm7
    vpaddb  %ymm6, %ymm7, %ymm6
    vpsadbw %ymm5, %ymm6, %ymm6
    vpaddq  %ymm6, %ymm4, %ymm4
    vpaddq  %ymm0, %ymm15, %ymm0
    addq    $-2, %rbx
    jne .LBB0_8
Run Code Online (Sandbox Code Playgroud)

分析

首先,我们可以看到GCC代码使用的popcntl指令速度非常快,至少对于标量操作来说是这样。

Clang 使用我的机器上的AVX-2 SIMD 指令集生成汇编代码。这就是 Clang 生成的程序与 GCC 相比如此快的原因:由于 SIMD 指令,它可以并行操作许多项目。

Numba 生成与 Clang 非常相似的代码。这并不奇怪,因为 Numba 基于 LLVM-Lite(以及 LLVM),而 Clang 也基于 LLVM。然而,在解释性能影响方面存在细微差别。事实上,Numba 汇编代码运行的项目比 Clang 对应项目多两倍。这可以通过计算指令数量看出vpsrlw(8 VS 4)。我不希望这会产生什么影响,因为 Clang 循环已经很好地展开了,而且进一步展开它的好处很小。实际上,这种更激进的展开是一种副作用。主要区别在于Numba 运行在 64 位整数上,而 C 代码运行在 32 位整数上!这就是 Clang 以不同方式展开循环并生成不同指令的原因。事实上,较小的整数会导致 Clang 生成一系列指令来转换不同大小的整数,效率较低。恕我直言,这是影响优化器的副作用,因为对较小项目的操作通常可用于生成更快的 SIMD 代码。在这种情况下,LLVM 生成的代码似乎不是最佳的:它使我机器上的端口 5(即洗牌/置换执行单元)饱和,而人们可以编写不使其饱和的代码(但这并不容易)。


更快的 C 实现

您可以修复 C++ 实现以便操作 64 位整数

#include <stdio.h>
#include <time.h>
#include <stdint.h>

// Function to calculate the population count (number of set bits) of an integer using __builtin_popcount
uint64_t popcount(uint64_t num) {
    return __builtin_popcountl(num);
}

int main() {
    int64_t n;
    printf("Enter the value of n: ");
    scanf("%ld", &n);

    // Variables to store start and end times
    struct timespec start_time, end_time;

    // Get the current time as the start time
    clock_gettime(CLOCK_MONOTONIC, &start_time);

    int64_t sum = 0;
    for (int64_t i = 0; i < n; i++) {
        sum += popcount(i);
    }

    // Get the current time as the end time
    clock_gettime(CLOCK_MONOTONIC, &end_time);

    // Calculate the elapsed time in microseconds
    long long elapsed_time = (end_time.tv_sec - start_time.tv_sec) * 1000000LL +
                            (end_time.tv_nsec - start_time.tv_nsec) / 1000;

    printf("Sum of population counts from 0 to %ld-1 is: %ld\n", n, sum);
    printf("Elapsed time: %lld microseconds\n", elapsed_time);

    return 0;
}
Run Code Online (Sandbox Code Playgroud)

这会在我的机器上使用 Clang 生成与 Numba 一样快的程序(GCC 仍然生成缓慢的标量实现)。


笔记

仅当您的实际代码是 SIMD 友好的(即popcount可以应用于多个连续项)时,SIMD 版本才有意义。否则,标量实现的结果可能会截然不同(事实上,三个编译器生成非常接近的代码,我希望它同样快)。

AVX-512 提供的SIMD 指令VPOPCNTDQ应该明显优于使用(仅)AVX-2 的 LLVM 生成的代码。由于我的机器上没有 AVX-512,并且 AVX-2 不提供这样的指令,因此它使得LLVM 的意义是使用 AVX-2 生成汇编代码。AVX-512 指令可以并行计算 16 x 32 位整数中 1 的数量,同时所花费的周期数与其标量对应指令大致相同。更准确地说,该指令只能使用指令集AVX512VPOPCNTDQ+来使用AVX512VL(AFAIK 并非在所有支持 AVX-512 的 CPU 上都可用)。截至目前,该指令仅适用于少数 x86-64 微架构(例如 Intel Ice-Lake、Intel Sapphire Rapids 和 AMD Zen4)。