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 特性可用于调用内部函数。 例如，让我们考虑 x86 和 x86_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