当我开始撰写答案时,我认为内联汇编是完全无意义的,而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;
}
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)部分是目标操作数的累加延迟;在源元素的水平对添加已经发生后,它允许来自先前相同类型操作的晚期转发。 因此,在使用
uadalp
的
sum += 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来说完全不好,除非你展开很多次循环。
cnt
结果单独缩减为标量,特别是一次只针对单个u64
而不是完整的128位。RustC是否会自动向量化u64.count_ones()
,如果您在数组元素上使用它?Clang会使用__builtin_popccountll
。或者说,如果自动向量化产生良好的汇编代码,则使用u8.count_ones()
。 - Peter Cordesaddv
是一种长延迟指令。如果您正在处理内存中的长数组,则应实现整个循环并使用addp
代替。 - Jake 'Alquimista' LEE