Module std::arch

1.27.0 · source ·
Expand description

SIMD 和供应商内部功能模块。

此模块旨在用作特定于体系结构的固有函数的门户,该固有函数通常与 SIMD 相关 (但并非总是如此!)。 Rust 编译到的每个体系结构都可能在此处包含一个子模块,这意味着这不是便携式模块! 如果您要编写可移植的库,请在使用这些 API 时多加注意!

在此模块下,您将找到一个以架构命名的模块,例如 x86_64。Rust 可以编译的每个 #[cfg(target_arch)] 此处可能都有一个模块条目,仅存在于该特定目标上。 例如,i686-pc-windows-msvc 目标在此处将具有 x86 模块,而 x86_64-pc-windows-msvc 具有 x86_64

Overview

该模块公开了特定于供应商的内部函数,这些内部函数通常对应于单个机器指令。 这些内部函数不是可移植的:它们的可用性取决于体系结构,并且并非该体系结构的所有机器都可以提供该内部函数。

arch 模块旨在作为高级 API 的实现细节。正确使用它可能会非常棘手,因为您需要确保至少遵守以下几点保证:

  • 使用了正确的体系结构模块。例如,arm 模块在 x86_64-unknown-linux-gnu 目标上不可用。 通常,通过在使用此模块时确保正确使用 #[cfg] 来完成此操作。
  • 程序当前正在运行的 CPU 支持被调用的函数。例如,在实际上不支持 AVX2 的 CPU 上调用 AVX2 函数是不安全的。

由于后者的保证,该模块中的所有内部函数都是 unsafe,因此在调用它们时要格外小心!

CPU 特性检测

为了以一种安全的方式调用这些 API,有许多机制可用来确保正确的 CPU 特性可用于调用内部函数。 例如,让我们考虑 x86x86_64 体系结构上的 _mm256_add_epi64 内部函数。 这个函数需要 AVX2 特性,由英特尔记录,所以为了正确调用这个函数,我们需要 (a) 保证我们仅在 x86/x86_64 和 (b) 上调用它,以确保 CPU 特性可用

静态 CPU 特性检测

我们可以使用的第一个选项是通过 #[cfg] 属性有条件地编译代码。CPU 特性对应于可用的 target_feature cfg,可以这样使用:

#[cfg(
    all(
        any(target_arch = "x86", target_arch = "x86_64"),
        target_feature = "avx2"
    )
)]
fn foo() {
    #[cfg(target_arch = "x86")]
    use std::arch::x86::_mm256_add_epi64;
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::_mm256_add_epi64;

    unsafe {
        _mm256_add_epi64(...);
    }
}
Run

在这里,我们使用 #[cfg(target_feature = "avx2")] 有条件地将此函数编译到我们的模块中。 这意味着,如果 avx2 特性是静态启用的,那么我们将在运行时使用 _mm256_add_epi64 函数。 可以通过使用 #[cfg] 来证明此处的 unsafe 块合理,仅在维护安全保证的情况下才编译代码。

静态启用一个特性通常是通过向编译器提供 -C target-feature-C target-cpu 标志来完成的。例如,如果您的本地 CPU 支持 AVX2,则可以使用以下命令编译上述函数:

$ RUSTFLAGS='-C target-cpu=native' cargo build

否则,您可以专门启用 AVX2 特性:

$ RUSTFLAGS='-C target-feature=+avx2' cargo build

请注意,在编译启用了特定特性的二进制文件时,确保仅在满足所需特性集的系统上运行二进制文件非常重要。

动态 CPU 特性检测

有时静态分派并不是您想要的。相反,您可能想构建一个可在各种 CPU 上运行的可移植二进制文件,但是在运行时它将选择可用的最优化的实现。 这使您可以构建 “最小公分母” 二进制文件,其中的某些部分针对不同的 CPU 进行了优化。

以之前的示例为例,我们将编译我们的二进制文件,而没有 AVX2 支持,但是我们只想为一个函数启用它。 我们可以按照以下方式进行操作:

fn foo() {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    {
        if is_x86_feature_detected!("avx2") {
            return unsafe { foo_avx2() };
        }
    }

    // fallback implementation without using AVX2
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
unsafe fn foo_avx2() {
    #[cfg(target_arch = "x86")]
    use std::arch::x86::_mm256_add_epi64;
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::_mm256_add_epi64;

    _mm256_add_epi64(...);
}
Run

这里有几个组件在起作用,所以让我们详细研究它们!

  • 首先,我们注意到 is_x86_feature_detected! 宏。 由标准库提供,此宏将执行必要的运行时检测,以确定程序所运行的 CPU 是否支持指定的特性。 在这种情况下,宏将扩展为一个布尔表达式,以评估本地 CPU 是否具有 AVX2 特性。

    请注意,与 arch 模块一样,此宏是特定于平台的。例如,在 ARM 上调用 is_x86_feature_detected!("avx2") 将是编译时错误。 为了确保我们不会遇到此错误,语句级别 #[cfg] 仅用于编译 x86/x86_64 上的宏用法。

  • 接下来,我们看到启用了 AVX2 的函数 foo_avx2。此函数用 #[target_feature] 属性修饰,该属性仅为此一个函数启用 CPU 特性。 使用 -C target-feature=+avx2 之类的编译器标志将为整个程序启用 AVX2,但使用属性将仅为一个函数启用它。 如此处所示,当前使用 #[target_feature] 属性要求函数也必须为 unsafe。 这是因为只能在具有 AVX2 的系统上正确调用该函数 (例如内部函数本身)。

有了所有这些,我们应该有一个有效的程序! 该程序将在所有计算机上运行,并且将在检测到支持的计算机上使用优化的 AVX2 实现。

Ergonomics

重要的是要注意,使用 arch 模块并不是世界上最简单的事情,因此,如果您想尝试一下,您可能会想方设法为自己做好准备!

该模块的主要目的是使 crates.io 上的稳定 crates 能够构建更多的人体工程学抽象,最终在引擎盖下使用 SIMD。 随着时间的流逝,这些抽象也可能会移入标准库本身,但是目前,此模块的任务是提供在稳定的 Rust 上使用供应商内部函数所需的最低限度的最低要求。

其他架构

本文档仅适用于一种特定的体系结构,您可以在以下位置找到其他文档:

Examples

首先,让我们看看实际上没有使用任何内部函数,而是使用 LLVM 的自动矢量化为 AVX2 和默认平台生成优化的矢量化代码。

fn main() {
    let mut dst = [0];
    add_quickly(&[1], &[2], &mut dst);
    assert_eq!(dst[0], 3);
}

fn add_quickly(a: &[u8], b: &[u8], c: &mut [u8]) {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    {
        // 请注意,此 `unsafe` 块是安全的,因为我们正在测试 `avx2` 特性确实在我们的 CPU 上可用。
        //
        if is_x86_feature_detected!("avx2") {
            return unsafe { add_quickly_avx2(a, b, c) };
        }
    }

    add_quickly_fallback(a, b, c)
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
unsafe fn add_quickly_avx2(a: &[u8], b: &[u8], c: &mut [u8]) {
    add_quickly_fallback(a, b, c) // 下面的函数内联在这里
}

fn add_quickly_fallback(a: &[u8], b: &[u8], c: &mut [u8]) {
    for ((a, b), c) in a.iter().zip(b).zip(c) {
        *c = *a + *b;
    }
}
Run

接下来,让我们来看一个手动使用内部函数的示例。在这里,我们将使用 SSE4.1 特性来实现十六进制编码。

fn main() {
    let mut dst = [0; 32];
    hex_encode(b"\x01\x02\x03", &mut dst);
    assert_eq!(&dst[..6], b"010203");

    let mut src = [0; 16];
    for i in 0..16 {
        src[i] = (i + 1) as u8;
    }
    hex_encode(&src, &mut dst);
    assert_eq!(&dst, b"0102030405060708090a0b0c0d0e0f10");
}

pub fn hex_encode(src: &[u8], dst: &mut [u8]) {
    let len = src.len().checked_mul(2).unwrap();
    assert!(dst.len() >= len);

    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    {
        if is_x86_feature_detected!("sse4.1") {
            return unsafe { hex_encode_sse41(src, dst) };
        }
    }

    hex_encode_fallback(src, dst)
}

// translated from
// <https://github.com/Matherunner/bin2hex-sse/blob/master/base16_sse4.cpp>
#[target_feature(enable = "sse4.1")]
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
unsafe fn hex_encode_sse41(mut src: &[u8], dst: &mut [u8]) {
    #[cfg(target_arch = "x86")]
    use std::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    let ascii_zero = _mm_set1_epi8(b'0' as i8);
    let nines = _mm_set1_epi8(9);
    let ascii_a = _mm_set1_epi8((b'a' - 9 - 1) as i8);
    let and4bits = _mm_set1_epi8(0xf);

    let mut i = 0_isize;
    while src.len() >= 16 {
        let invec = _mm_loadu_si128(src.as_ptr() as *const _);

        let masked1 = _mm_and_si128(invec, and4bits);
        let masked2 = _mm_and_si128(_mm_srli_epi64(invec, 4), and4bits);

        // return 0xff corresponding to the elements > 9, or 0x00 otherwise
        let cmpmask1 = _mm_cmpgt_epi8(masked1, nines);
        let cmpmask2 = _mm_cmpgt_epi8(masked2, nines);

        // add '0' or the offset depending on the masks
        let masked1 = _mm_add_epi8(
            masked1,
            _mm_blendv_epi8(ascii_zero, ascii_a, cmpmask1),
        );
        let masked2 = _mm_add_epi8(
            masked2,
            _mm_blendv_epi8(ascii_zero, ascii_a, cmpmask2),
        );

        // interleave masked1 and masked2 bytes
        let res1 = _mm_unpacklo_epi8(masked2, masked1);
        let res2 = _mm_unpackhi_epi8(masked2, masked1);

        _mm_storeu_si128(dst.as_mut_ptr().offset(i * 2) as *mut _, res1);
        _mm_storeu_si128(
            dst.as_mut_ptr().offset(i * 2 + 16) as *mut _,
            res2,
        );
        src = &src[16..];
        i += 16;
    }

    let i = i as usize;
    hex_encode_fallback(src, &mut dst[i * 2..]);
}

fn hex_encode_fallback(src: &[u8], dst: &mut [u8]) {
    fn hex(byte: u8) -> u8 {
        static TABLE: &[u8] = b"0123456789abcdef";
        TABLE[byte as usize]
    }

    for (byte, slots) in src.iter().zip(dst.chunks_mut(2)) {
        slots[0] = hex((*byte >> 4) & 0xf);
        slots[1] = hex(*byte & 0xf);
    }
}
Run

Re-exports

  • pub use core::arch::*;

Macros