如何优化AVX512下u8转u64的多值SIMD复制扩展（Splat/Broadcast）

我需要将若干u8类型数据扩展为u64类型，但并非直接支持的零扩展或符号扩展，而是执行「复制扩展」操作——即把每个输入字节重复8次，填充到对应的u64结果中。比如输入字节0xAB，对应的输出u64就是0xABABABABABABABAB。

示例Rust代码 #

#![feature(portable_simd)]

use std::simd::*;

// 逐个字节复制扩展的标量实现
pub fn batch_splat_scalar(x: [u8; 16]) -> [u64; 16] {
    let mut ret = [0; 16];
    for i in 0..16 {
        ret[i] = u64::from_le_bytes([x[i], x[i], x[i], x[i], x[i], x[i], x[i], x[i]]);
    }
    ret
}

// 基于标量实现的SIMD包装
pub fn batch_splat_simd(x: u8x16) -> u64x16 {
    Simd::from_array(batch_splat_scalar(x.to_array()))
}

当前编译器生成的汇编分析 #

上述代码在AVX512环境下编译后，编译器采用了零扩展+乘法+移位+或的组合方式实现：

1. 用vpmovzxbq将8个u8零扩展为8个u64；
2. 通过与预定义的掩码常量相乘，分别生成低32位和高32位的重复字节；
3. 移位合并高低位，最终得到每个字节重复8次的u64结果。

对应的汇编代码如下：

vpmovzxbq       zmm0, qword ptr [rsi]
        vpbroadcastq    zmm1, qword ptr [rip + .LCPI0_0]
        mov     rax, rdi
        vpmuludq        zmm2, zmm0, zmm1
        vpbroadcastq    zmm3, qword ptr [rip + .LCPI0_1]
        vpmuludq        zmm0, zmm0, zmm3
        vpsllq  zmm0, zmm0, 32
        vporq   zmm0, zmm2, zmm0
        vmovdqu64       zmmword ptr [rdi], zmm0
        vpmovzxbq       zmm0, qword ptr [rsi + 8]
        vpmuludq        zmm1, zmm0, zmm1
        vpmuludq        zmm0, zmm0, zmm3
        vpsllq  zmm0, zmm0, 32
        vporq   zmm0, zmm1, zmm0
        vmovdqu64       zmmword ptr [rdi + 64], zmm0
        vzeroupper
        ret

AVX512最优实现方案 #

AVX512BW指令集提供了**vpbroadcastb**指令，直接支持将每个输入字节广播到对应64位元素的所有8个字节中，完美匹配我们的需求，仅需3条核心指令即可完成16个u8到u64的复制扩展：

最优汇编序列

vmovdqu8 zmm0, [rsi]       ; 加载16个u8到zmm寄存器低16字节
vpbroadcastb zmm1, zmm0    ; 将每个u8广播到对应64位元素的8个字节（完成复制扩展）
vmovdqu64 [rdi], zmm1      ; 存储16个u64结果
vzeroupper
ret

对应的Rust内联汇编实现

如果需要手动控制指令序列，可以用内联汇编实现：

use std::simd::{u8x16, u64x16};

pub fn batch_splat_avx512(x: u8x16) -> u64x16 {
    let mut result = u64x16::default();
    unsafe {
        core::arch::asm!(
            "vmovdqu8 zmm0, {input}",
            "vpbroadcastb zmm1, zmm0",
            "vmovdqu64 {output}, zmm1",
            input = in(reg) x.as_ptr(),
            output = out(reg) result.as_mut_ptr(),
            out("zmm0") _,
            out("zmm1") _,
        );
    }
    result
}

方案优势 #

• 指令数更少：从原有的10+条指令缩减到3条核心指令，大幅降低指令延迟和吞吐量开销；
• 无额外常量依赖：不需要预定义乘法掩码，减少数据缓存占用；
• 直接匹配需求：vpbroadcastb原生支持字节到64位元素的全复制扩展，无需通过乘法/移位模拟。

内容的提问来源于stack exchange，提问作者ajp