如何最安全地读取数组末尾以启用 simd 矢量化?

问题描述 投票:0回答:1

我有一个长度正好为 15 的数组。如果它的长度正好为 16,我的函数就会有一个很好的优化 simd 实现。出于性能原因,我想假装它的长度为 16,并从数组末尾读取一个字节(这样我的函数结果就不会受到该字节内容的影响)。

有安全的方法吗?我最不坏的选择是什么?

#![feature(portable_simd)]
#![allow(invalid_reference_casting)]

pub fn find_first_16(
    x: u8,
    y: &[u8; 16],
) -> Option<usize> {
    use std::simd::cmp::SimdPartialEq;
    use std::simd::u8x16;
    let x = u8x16::splat(x);
    let y = u8x16::from_array(*y);
    let e = x.simd_eq(y);
    e.first_set()
}

// SAFETY: y must point to at least 16 bytes of mapped
// memory, though the bytes that are not part of the
// slice do not necessarily need to be initialized,
// part of the same allocation, or allocated at all.
pub unsafe fn find_first_assume_16(
    x: u8,
    y: &[u8]
) -> Option<usize> {
    let y = unsafe {
        &*(y as *const [u8] as *const [u8;16])
    };
    match find_first_16(x, y) {
        Some(ret) if ret < y.len() => Some(ret),
        _ => None
    }
}


#[repr(C)]
pub struct Foo([u8;15], u8);

pub fn find_first_foo(
    x: u8,
    y: &Foo
) -> Option<usize> {
    // SAFETY: Foo has a 16th byte which is part of
    // the same allocation, and is initialized.
    unsafe { find_first_assume_16(x, &y.0) }
}

#[repr(align(16))]
pub struct Bar([u8;15]);

pub fn find_first_bar(
    x: u8,
    y: &Bar
) -> Option<usize> {
    // SAFETY: Bar has a 16th byte which is part of the
    // same allocation, although it is not initialized.
    unsafe { find_first_assume_16(x, &y.0) }
}

pub fn find_first_yolo(
    x: u8,
    y: &[u8;15]
) -> Option<usize> {
    // SAFETY: it's highly unlikely that `y` occupies the
    // final 15 bytes of a mapped region of memory. Even
    // if it does, we'll get a segfault, which is safe.
    unsafe { find_first_assume_16(x, y) }
}

注意 find_first_{foo,bar,yolo} 向量化得很好

        vmovd   xmm0, edi
        vpbroadcastb    xmm0, xmm0
        vpcmpeqb        xmm0, xmm0, xmmword ptr [rsi]
        vpmovmskb       eax, xmm0
        test    eax, eax
        setne   cl
        rep       bsf edx, eax
        cmp     dx, 15
        setb    al
        and     al, cl
        movzx   eax, al
        ret

https://godbolt.org/z/fj168hrPa

rust vectorization simd
1个回答
0
投票

这在 Rust 中实际上是不可能的,因为它是未定义的行为。例如,仅因为

Bar
是 16 字节对齐,并不意味着读取 16 字节数量是安全的。
Bar
可以位于分配有
mmap
的区域中 16 字节倍数的地址,但最后一个字节未映射到文件中(因此理论上可以是 SIGBUS)。

你会遇到 UB 问题的部分原因是额外的字节可以被其他代码使用,虽然你只有一个对 15 字节数据的不可变引用,但另一个线程可能有一个可变引用到额外的字节并同时修改它。这是一场数据竞争,Rust 不允许这种情况发生。

即使您从未编写过这样的数据竞争,编译器也可以根据所编写的代码进行优化,根据通常所说的 as-if 规则:编译器实际上可以发出它想要的任何代码,只要代码的功能就好像它是一个忠实的翻译。因为不允许您编写实现未定义行为的代码,所以编译器可以假设它不会发生,因此如果您确实尝试这样做,您的代码最终可能会出现奇怪的行为不当。 (然而,与 C 不同,Rust 通常会尝试确保安全代码永远不会触发未定义的行为,因此它没有 C 那样的巨大地雷。)

© www.soinside.com 2019 - 2024. All rights reserved.