高效计算汉明重量。

3

我使用的是苹果M1处理器。

我的目标是高效地在Rust中计算一个大型char数组中的所有1位。我查阅了ARM NEON指令,我认为可以通过cnt指令(按8位块计算1)然后再执行addv来添加每个8位块来完成该操作。

开始时,我打算输入一个64位无符号整数。

fn aarch64_hamming_weight(inp: u64) -> u64 {
    let mut outp: u64;
    unsafe {
        asm!(
            "dup v0.2d, {input}",
            "cnt v0.16b, v0.16b",
            "addv b1, v0.16b",// this adds all the bytes of the vector, how do i get the result?!
            //"fmov {output2}, b1",
            "fmov {output}, v0.d[1]",
            input = in(reg) inp,
            output = out(reg) outp,
        )
    }

    return outp;
}

这几乎可以工作,但如果你的数字大于8位,则结果不完全正确;225的二进制计数为4,425的二进制计数为260。
所以, 1. 我需要一种获取addv输出的方法。 2. 我需要一种使其能够与任意字符数组一起使用的方法(这需要一个循环,每次提取128位)。

1
如果你想计算整个数组,请不要将每个cnt结果单独缩减为标量,特别是一次只针对单个u64而不是完整的128位。RustC是否会自动向量化u64.count_ones(),如果您在数组元素上使用它?Clang会使用__builtin_popccountll。或者说,如果自动向量化产生良好的汇编代码,则使用u8.count_ones() - Peter Cordes
2
顺便提一下,https://github.com/WojciechMula/sse-popcount/ 上有一些ARM Neon手动向量化的代码。如果它像同一存储库中的x86代码一样经过良好调整,那么你应该以此为目标,无论你如何实现,都要得到最终的汇编输出。 - Peter Cordes
2
addv 是一种长延迟指令。如果您正在处理内存中的长数组,则应实现整个循环并使用 addp 代替。 - Jake 'Alquimista' LEE
数组有多大?总位数能否适应u16、u32、u64?这对算法并没有太大影响,但如果它适合16或32位,你可以在每次迭代中节省一两个“addp”。 - Nate Eldredge
输入数组大约有32MB,它是一个具有2.5亿位的位集合,我正在尝试计算设置成员。 - Anko
1个回答

4
当我开始撰写答案时,我认为内联汇编是完全无意义的,而RustC将会很好地对u8.count_ones()进行向量化。不幸的是情况并非如此,但如果您要编写整个循环而不是一个u64一次,则应仅使用asm(可能仍然是内联汇编,但编写整个函数可能是合理的)。如果Rust具有ARM / AArch64的SIMD内置函数,那将比内联汇编更好,并且绝对值得使用。
为了计算整个数组,您不希望将每个cnt结果单独缩减到标量,特别是不要一直缩减到通用整数寄存器。特别是针对一个u64而不是完整的128位。将其添加到仅在最后缩减的计数向量中。
clang/LLVM知道如何自动向量化C语言中的popcount __builtin_popcount(),所以我希望基于LLVM的RustC对AArch64也能做得很好。对于u64来说还可以,但对于u8来说非常糟糕。如果您可以安全地将u64跨度指向数据(或引用或无论Rust如何处理?),则可以在可移植性良好且没有脆弱的内联汇编的情况下获得良好的效果。并且希望未来的编译器版本可以改进。
pub fn hamming_weight_arr(inp: &[u64]) -> u64 {
    let mut sum : u64 = 0;
    for v in inp {
        sum += v.count_ones() as u64;  // maybe not terrible but a lot of hsumming.
        // doing ldp q2, q3 to load 32 bytes per iteration
    }
    return sum;
}

使用-O --target aarch64-unknown-linux-gnu编译,链接为https://godbolt.org/z/7qP7cK6bv(默认使用-C --opt-level=3)。

... some setup
.LBB1_5:                                 // inner loop
        ldp     q2, q3, [x8, #-16]       // load pair of 16-byte vectors = unroll by 4x u64
        add     x8, x8, #32              // pointer increment by 32 bytes
        subs    x12, x12, #4
        cnt     v2.16b, v2.16b
        cnt     v3.16b, v3.16b
        uaddlp  v2.8h, v2.16b            // hsum byte pairs to u16 halfwords
        uaddlp  v3.8h, v3.16b
        uaddlp  v2.4s, v2.8h             // hsum u16 pairs to u32 words
        uaddlp  v3.4s, v3.8h
        uadalp  v0.2d, v2.4s          // sum u32 pairs into accumulator u64 vectors (doublewords)
        uadalp  v1.2d, v3.4s
        b.ne    .LBB1_5
        add     v0.2d, v1.2d, v0.2d        // combine pair of accumulators
        cmp     x10, x11
        addp    d0, v0.2d                  // and reduce to one 64-bit
        fmov    x8, d0
        b.eq    .LBB1_9
.LBB1_7:
        add     x10, x0, x1, lsl #3
.LBB1_8:
        ldr     d0, [x9], #8               // scalar cleanup loop, one u64 at a time
        cmp     x9, x10
        cnt     v0.8b, v0.8b
        uaddlv  h0, v0.8b                  // slow instruction, as Jake mentioned.  Or at least high latency
        fmov    w11, s0
        add     x8, x11, x8
        b.ne    .LBB1_8
.LBB1_9:
        mov     x0, x8
        ret

您可能认为使用sum: u32可以帮助减少循环内部的扩展。 (如果您有可能溢出的大型数组,请使用外部循环)。但实际上,RustC仍然会扩展到64位,但然后会更多地工作以将这些计数截断为32位。几乎肯定是一个未优化的问题。(也许是围绕x86 psadbw构建的策略,它可以在一步中将字节求和为u64块; LLVM使用pshufb自动矢量化popcount,该矢量化适用于x86的AVX2)。
您可能认为对于u8数组做同样的事情应该矢量化为相同的代码,只需要一些额外的标量清理,对吧?好吧,它应该,但实际上它仍然只展开了4个,就像LLVM喜欢做的那样,但是内部循环变成了垃圾。
// &[u8] version  inner loop is a disaster

LBB2_5:
        ldurb   w12, [x8, #-2]            // scalar zero-extending byte load
        subs    x11, x11, #4
        ldrb    w13, [x8]                 // scalar sign-extending byte load
        fmov    d2, x12                   // copy it into vector reg
        ldurb   w12, [x8, #-1]
        fmov    d3, x13
        ldrb    w13, [x8, #1]
        add     x8, x8, #4
        mov     v2.d[1], x12              // get two more bytes into the high 64 bits of a vector
        mov     v3.d[1], x13
        cnt     v2.16b, v2.16b           // same cnt / uaddlp sequence as before
        cnt     v3.16b, v3.16b
        uaddlp  v2.8h, v2.16b
        uaddlp  v3.8h, v3.16b
        uaddlp  v2.4s, v2.8h
        uaddlp  v3.4s, v3.8h
        uadalp  v0.2d, v2.4s
        uadalp  v1.2d, v3.4s
        b.ne    .LBB2_5

所以它正在将 (v as u64).count() 转换为向量,并使用预设的方案进行优化。例如,如果除了每个 64 位块的低字节之外,cnt 结果都是零,则使用 uaddlp 扩展是没有意义的,他们可以只使用带有 64 位块的垂直 add
与您从https://github.com/WojciechMula/sse-popcount/手动向量化的ARM代码得到的C编译器结果进行比较。如果该ARM Neon代码像同一存储库中的x86代码一样经过优化,那么您应该以此为目标来生成最终的汇编输出,无论您如何实现。

我猜想,仅将每字节计数扩展到16位的内部循环将是不错的选择,运行次数为可能不会因为在进入其中的一对字节中看到全1而导致溢出的+=16。即65535/16向下取整=4095个向量,然后扩展到64位块。

Mula的版本仅进行31个内部循环迭代,使用vaddq_u8累加到8位元素中。但是uadalp v0.16b, v2.8h将u8字节水平成对累加为u16半字元素,在32位NEON中不可用,只有AArch64 ASIMD可用。

关键点是在最内层循环中尽可能少地进行工作,理想情况下每个cnt结果向量仅使用一条廉价指令。如果您可以在累加时免费获得一些扩展(或者廉价到不会成为瓶颈),那么执行时间就可以更长而不会发生溢出。(并且意味着外部循环中后续的水平约简步骤更便宜。)

uadalp与纯竖直的add在某些CPU上的性能略有不同,但很可能值得使用。 Cortex-A57优化指南指出它具有4(1)周期延迟,1 /时钟吞吐量。 (1)部分是目标操作数的累加延迟;在源元素的水平对添加已经发生后,它允许来自先前相同类型操作的晚期转发。 因此,在使用uadalpsum += pairs循环中,循环传递的依赖链仅长1个周期。 这是理想的。

整数向量上的常规add具有3个周期延迟,2 /时钟吞吐量,因此它具有更好的吞吐量,但会创建一个3个周期的循环传递依赖链。 (而且不会完成一步水平累加工作,并且由于仅使用8位和,所以会更快溢出。)

在Cortex-A57上,cnt的吞吐量也只有每个时钟周期1次,因此uadalp的每个时钟周期吞吐量不是总体瓶颈,除非它们争用同一个执行端口。 uadalp运行在F1上,SIMD整数add运行在F0 / F1上,cnt运行在F0 / F1上。 因此,无论如何,加法操作都需要窃取一些cnt吞吐量,问题在于当F1排队了大量未来的uadalp操作时,是否可以有效地将cnt调度到主要是端口F0,(在数据尚未准备好时;乱序执行向前看。 在大多数CPU上,操作被安排到端口中,因为它们从前端重命名/分配到后端。 CPU不知道数据将以什么顺序准备好,但它可以看到不同端口的队列长度。

这是可能的(伪代码)

// Probably not a good idea
c1 = cnt(load1)
c2 = cnt(load2)
c1 += c2    // plain SIMD vertical add
uadalp accumulator, c1

这意味着在循环中只有一个,以防止它成为吞吐量瓶颈,同时仍然只将用于累加器,通过累加器保持循环依赖链的长度较短。(假设其他AArch64 CPU也对积累指令进行早期转发)。这使得一次迭代内的短独立依赖链略微变长(从加载到输入累加),而不是保持为两个单独的依赖链。低端ARM CPU通常具有非常有限的乱序执行能力(小的调度程序缓冲区),以找到指令级并行性并在循环迭代之间重叠工作。保持非循环依赖链短,可以更轻松地让CPU实现这一点。但这对于类似Cortex-A53的顺序执行的AArch64来说完全不好,除非你展开很多次循环。

1
这需要我一些时间来理解,因为我对很多术语仍然很陌生,但你给了我很好的指引。 1)将向量添加到其他向量比逐对添加更有效,因此请暂时不要这样做。 2)保持向量为8位更有效,只有在必要时才进行扩展。 3)使用纯Rust编写一些代码并查看汇编输出是学习和入门的好方法 :) 非常感谢。 - Anko
1
@Anko:如果有一个uadalp v0.8h,v2.16b,那实际上是将字节对相加成为u16元素的向量的理想选择,假设它的运行速度与add v0.16b,...一样快。根据https://developer.arm.com/documentation/100069/0606/SIMD-Vector-Instructions/UADALP--vector-,确实存在这种指令。所以你不需要将它们保留为8位,你只需要在最内层循环中尽可能少地进行操作。如果你可以在累加时免费进行扩展,那么执行就可以在内部循环中更长时间地停留,而不会出现溢出的可能性。 - Peter Cordes
1
@Anko:顺便说一句,在好奇4(1)对于Cortex-A57上的uadalp意味着什么延迟之后,我更新了我的答案。实际上,它比我想象的要好得多,只允许前一个累加操作进行“晚转发”,以解决循环传递依赖性问题,只需要1个周期的延迟。 - Peter Cordes
非常感谢你,彼得,你真的让我的一天。我正在研究aarch64/simd的rust内在状态,并尝试使用各种汇编实现。你给我的那个链接网站(godbolt)简直太棒了! - Anko
1
@Anko:完全同意Godbolt非常方便,查看编译器输出以了解正在发生的事情对于调整高性能循环至关重要。(并且为了学习如何在给定的ISA上进行汇编操作。有关编写编译为有趣汇编的简单函数的更多信息,请参见如何从GCC / clang汇编输出中删除“噪音”?,例如获取参数并返回值而不是对常量进行操作。) - Peter Cordes
1
这种只使用内置函数的技术比标准的逐字节 Rust 方式快了 14.73 倍。非常感谢 : )。我认为可能还有进一步的优化,但我已经转向问题的另一个部分。对于超过 192.00 MB 的计数,只需 3.67 毫秒即可完成 :)。 - Anko

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接