Dance with Compiler - EP3 ARM64 汇编传参约定以及 restrict 汇编分析

在 ARM64 架构（也称为 AArch64）中，函数调用约定定义了寄存器如何用于传递参数和返回值。这些约定有助于实现高效的函数调用和返回。在 ARM64 的汇编中，寄存器传参遵循以下约定：

参数传递寄存器

x0 - x7: 这 8 个寄存器用于传递函数的前 8 个参数（对于整数类型的参数）。如果函数有更多的参数，这些额外的参数会通过栈传递。

x0 用于第一个参数
x1 用于第二个参数
x2 用于第三个参数
x3 用于第四个参数
x4 用于第五个参数
x5 用于第六个参数
x6 用于第七个参数
x7 用于第八个参数

v0 - v7: 对于浮点数和 SIMD（单指令多数据）参数，这 8 个寄存器用于传递前 8 个浮点数参数（每个寄存器可以存储一个 64 位的浮点数，或者一组 128 位的 SIMD 数据）。

返回值

x0: 用于返回整数类型的结果（例如，int、long、pointer 等）。

v0: 用于返回浮点数或 SIMD 类型的结果（例如，float、double、__m128 等）。

阅读汇编代码时需要注意上述约定。

学完后实习一下：https://godbolt.org/z/zWjbo9183

仅仅是用了 restrict，性能可以提升非常大：

source：here

C++代码：

struct AvgState {
  uint64_t numerator{0};
  uint64_t denominator{0};

  double divide() { return numerator / denominator; }
};


void addBatch(size_t batch_size, AggregateDataPtr __restrict state, Column *args) __attribute__((noinline)) {
    for (size_t i = 0; i < batch_size; ++i) {
     data(state).numerator += args[0].data()[i];
        ++(data(state).denominator);

    }
}

对应汇编代码：

asm 复制代码

IAggregate<AvgAggregator>::addBatch(unsigned long, char*, Column*) [clone .isra.0]:
        cbz     x0, .L7                 <= x0 是第一个参数，batch_size，cbz（compare and branch on zero），batch_size 是 0 的话直接返回
        ldr     x2, [x2]                <= x2 是第三个参数 args，args[0] 地址等于 x2, args[0].data() 地址也等于 x2, args, args[0].data()[0]地址也等于 x2
        ldp     x3, x6, [x1]            <= x1 指向了 state，state 结构包含了两个地址相连的成员 numerator 和 denominator，所以这个指令获得了他们的地址：x3 = state.numerator.  x6 = state.denominator
        add     x5, x2, x0, lsl 3     <= args[0].data()[0] + (batch_size << 3) 得到 args[0].data() 最后一个元素的地址，用于控制循环结束。也就是说，这里是通过 address guard 的方式来控制循环结束
.L9:
        ldr     x4, [x2], 8           <= 将 x2 里的内容载入 x4，同时将 x2 加上 8（专门针对 for 循环场景设计的指令）。一条指令实现了两个能力：args[0].data()[i] 取值，i++
        add     x3, x3, x4         <= x3 = state.numerator，x4 = args[0].data()[i] ，这条指令计算 state.numerator + args[0].data()[i] 
        cmp     x5, x2           <= 判断循环是否结束（x2 的地址是否抵达了上面计算的边界）
        bne     .L9                <= 如果还没有到边界，则跳到 L9 继续循环
        add     x0, x0, x6      <= 我们可以发现，循环里没有执行过 ++(data(state).denominator); 操作。这里一把梭哈，(data(state).denominator) = (data(state).denominator) + batch_size  减少了很多指令执行。
        stp     x3, x0, [x1].      <= 把 state.numerator， state.denominator 的最新值写会到。state 结构的内存里，完成全部计算
.L7:
        ret

这些汇编代码非常简洁，从性能角度，最最重要的一点是循环中只有一处 ldr 操作，其余都是寄存器里的算术运算。我们知道，在性能领域，访问内存往往是瓶颈所在。

但是，如果 state 上没有加 __restrict，则是完全另一幅光景：

IAggregate<AvgAggregator>::addBatchWithoutOpt(unsigned long, char*, Column*) [clone .isra.0]:
        cbz     x0, .L1  <= x0 指向 batch_size
        ldr     x5, [x2] <= x5 指向 args[0].data()[0]
        ldp     x3, x2, [x1] <= x1 指向 state，x3 = state.numerator.  x2 = state.denominator
        add     x4, x0, x2  <= x4 =  batch_size + state.denominator，也就是用 state.denominator 终值做循环结束条件
        sub     x5, x5, x2, lsl 3 <= x5  指向 args[0].data() 最后一个元素的地址，用于控制循环结束
.L3:
        ldr     x0, [x5, x2, lsl 3] <= args[0].data()[0] + x2 << 3，也就是访问  args[0].data()[i]
        add     x2, x2, 1. <= 循环加1, state.denominator++
        add     x3, x3, x0 <= args[0].data()[i]  + state.numerator -> state.numerator
        stp     x3, x2, [x1].          <= 将 x3, x2 的内容写回 x1 地址。即更新 state 在内存中的值
        cmp     x2, x4    <= 判断是否已经循环结束
        bne     .L3
.L1:
        ret

在上面的循环里，除了 ldr 访存，还多了一个 stp 写内存操作。二者巨大的性能差异也是因为这条指令而起。

为什么没有 __restrict 后性能差异如此巨大呢？我们来分析下函数签名背后蕴含的可能：

void addBatch(size_t batch_size, AggregateDataPtr __restrict state, Column *args) __attribute__((noinline)) {

state 没有使用 restrict 时，编译器必须假设 strict 可能指向了 args。如果 state 指向了 args，那么我们看循环里的两条语句可能发生什么情况：

cpp 复制代码

for  (size_t i = 0; i < batch_size; ++i) {
     data(state).numerator += args[0].data()[i];   <= numerator 被更新，意味着 args[0].data 指针本身，args[0].data 里的元素，都可能被更新。编译器无法判断这是否可能，必须做最坏打算。并且，如果真的是这样，编译器还必须假设这是用于有意为之，它必须保证用户能得到符合预期的结果。
     ++(data(state).denominator); <= 这一步也是和上面一样，denominator 被更新，也意味着有一段内存被更新了，这段内存是什么？不知道，不能做任何假设。
     // 所以，到这里的时候，编译器必须把对 state 的更新写回到内存，只有这样，下一次循环才能得到"符合预期"的行为。
}

更凶猛的后果

上面是用 -O2 编译的，如果使用 -O3，还可以看到更凶猛的结果。下面分别展示了 x86 上 -O3 编译和 arm64 的 -O3 编译结果：

X86：

asm; 复制代码

IAggregate<AvgAggregator>::addBatchWithoutOpt(unsigned long, char*, Column*) [clone .constprop.0]:
        mov     rax, QWORD PTR [rdi+8]
        mov     r9, rsi
        mov     rdx, QWORD PTR [rdi]
        mov     rcx, QWORD PTR [r9]
        mov     r8, rax
        lea     rsi, [rax+1048576]
        neg     r8
        lea     rcx, [rcx+r8*8]
.L2:
        add     rdx, QWORD PTR [rcx+rax*8]
        add     rax, 1
        mov     QWORD PTR [rdi], rdx
        mov     QWORD PTR [rdi+8], rax
        cmp     rax, rsi
        jne     .L2
        ret
IAggregate<AvgAggregator>::addBatch(unsigned long, char*, Column*) [clone .constprop.0]:
        mov     r8, rsi
        mov     rcx, QWORD PTR [rdi+8]
        mov     rsi, QWORD PTR [rdi]
        pxor    xmm0, xmm0
        mov     rax, QWORD PTR [r8]
        lea     rdx, [rax+8388608]
.L6:
        movdqu  xmm2, XMMWORD PTR [rax]
        add     rax, 16
        paddq   xmm0, xmm2
        cmp     rdx, rax
        jne     .L6
        movdqa  xmm1, xmm0
        add     rcx, 1048576
        psrldq  xmm1, 8
        mov     QWORD PTR [rdi+8], rcx
        paddq   xmm0, xmm1
        movq    rax, xmm0
        add     rax, rsi
        mov     QWORD PTR [rdi], rax
        ret

ARM64：

IAggregate<AvgAggregator>::addBatchWithoutOpt(unsigned long, char*, Column*) [clone .constprop.0]:
        ldr     x4, [x1]
        ldr     x1, [x0, 8]
        ldr     x2, [x0]
        add     x5, x1, 1048576
        sub     x4, x4, x1, lsl 3
.L2:
        ldr     x3, [x4, x1, lsl 3]
        add     x1, x1, 1
        add     x2, x2, x3
        stp     x2, x1, [x0]
        cmp     x1, x5
        bne     .L2
        ret
IAggregate<AvgAggregator>::addBatch(unsigned long, char*, Column*) [clone .constprop.0]:
        ldr     x1, [x1]
        ldp     x4, x3, [x0]
        add     x2, x1, 8388608
        movi    v0.4s, 0
.L6:
        ldr     q1, [x1], 16
        add     v0.2d, v0.2d, v1.2d
        cmp     x2, x1
        bne     .L6
        addp    d0, v0.2d
        add     x1, x3, 1048576
        str     x1, [x0, 8]
        fmov    x1, d0
        add     x1, x1, x4
        str     x1, [x0]
        ret

可以看到，此时用到了 SIMD 指令：
movi v0.4s, 0, add v0.2d, v0.2d, v1.2d, addp d0, v0.2d, fmov x1, d0，一定程度上可以加速执行。

不过也需要注意，上面 2d 表示一条指令只能同时处理两个 64 位整数，也许快不了太多。得256、512 bit 的 SIMD 才能更显神威。

BTW，这篇文章讲 SIMD 以及汇编指令讲得挺不错：https://no5-aaron-wu.github.io/2022/06/14/SIMD-3-NeonAssembly/