Press "Enter" to skip to content

为了加速你的Rust代码的SIMD,需要遵循九个规则(第一部分)

通过对range-set-blaze箱的数据摄取提速7倍的一般经验

将计算委托给小蟹的螃蟹 — 来源:https://openai.com/dall-e-2/。其他所有图片均来自作者。

感谢西雅图 Rust Meetup 的 Ben Lichtman (B3NNY) 指引我在SIMD方向上的正确前进。

SIMD(单一指令,多个数据)操作自20世纪初以来就是英特尔/AMD和ARM CPU的特性。这些操作使您能够仅使用单个核心上的一个CPU操作就将8个<i32数组添加到另一个8个<i32数组中。使用SIMD操作极大地加快了某些任务的速度。如果您不使用SIMD,可能没有充分利用您的CPU能力。</i32</i32

这是一个“又一个 Rust 和 SIMD”文章吗?是和不是。是的,我确实将SIMD应用于一个编程问题,然后感觉有必要写一篇文章。不,我希望本文也能讲得足够深入,以指引您完成项目。它解释了 Rust 的最新可用的 SIMD 能力和设置。它包含了一个 Rust SIMD 的小抄。它展示了如何在不离开安全 Rust 的情况下将您的 SIMD 代码变成通用代码。它启动了一些使这一过程更加简便的新 cargo 命令。

range-set-blaze 箱使用 RangeSetBlaze::from_iter 方法将可能很长的整数序列摄取进来。当整数是“聚团”的时候,它可以比 Rust 的标准 HashSet::from_iter 方法运行 30倍快。如果我们使用 SIMD 操作,我们还能做得更好吗?是的!

关于“聚团”的定义,请参见前一篇文章的第二条规则。另外,如果整数不是聚团的,那会发生什么?RangeSetBlazeHashSet2到3倍

对于聚团的整数,基于 SIMD 操作的 RangeSetBlaze::from_slice 新方法比 RangeSetBlaze::from_iter 快7倍。这使得它比 HashSet::from_iter 快200倍以上。(当整数不是聚团的时候,它仍然比 HashSet 慢2到3倍。)

在实现这个加速过程中,我学到了九个规则,可以帮助您通过 SIMD 操作加速您的项目。

这些规则是:

  1. 使用 Rust 的 nightly 版本和 core::simd,Rust 实验中的标准 SIMD 模块。
  2. CCC:检查、控制和选择计算机的 SIMD 能力。
  3. 有选择地了解 core::simd
  4. 构思候选算法。
  5. 使用 Godbolt 和 AI 来理解您的代码的汇编,即使您不懂汇编语言。
  6. 使用内联泛型(如果内联泛型不起作用,则使用)宏,以及(如果宏不起作用)特性,将您的 SIMD 代码泛化到所有类型和 LANES。

请参阅即将发布的第二部分,了解这些规则:

7. 使用 Criterion 基准测试来选择算法,并发现 LANES 应该(几乎)总是 32 或 64。

8.将您最好的SIMD算法与as_simdi128/u128的专用代码以及额外的上下文基准测试集成到项目中。

9.使用可选的货物特性将您最好的SIMD算法从项目中剔除(暂时)。

旁注:为了避免模棱两可,我将这些称为“规则”,但它们当然只是建议。

规则1:使用夜间版Rust和core::simd,Rust的实验性标准SIMD模块。

在尝试在更大的项目中使用SIMD操作之前,让我们确保我们能够让它们正常工作。以下是步骤:

首先,创建一个名为simd_hello的项目:

cargo new simd_hellocd simd_hello

编辑src/main.rs以包含(Rust playground):

// 告诉夜间版Rust启用'portable_simd'#![feature(portable_simd)]use core::simd::prelude::*;// 常量Simd结构const LANES: usize = 32;const THIRTEENS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([13; LANES]);const TWENTYSIXS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([26; LANES]);const ZEES: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([b'Z'; LANES]);fn main() {    // 从LANES字节的切片中创建Simd结构    let mut data = Simd::<u8, LANES>::from_slice(b"URYYBJBEYQVQBUBCRVGFNYYTBVATJRYY");    data += THIRTEENS; // 每个字节都加上13    // 将每个字节与'Z'进行比较,如果字节大于'Z',则减去26    let mask = data.simd_gt(ZEES); // 比较每个字节与'Z'    data = mask.select(data - TWENTYSIXS, data);    let output = String::from_utf8_lossy(data.as_array());    assert_eq!(output, "HELLOWORLDIDOHOPEITSALLGOINGWELL");    println!("{}", output);}

接下来 – 完整的SIMD功能需要Rust的夜间版。假设您已经安装了Rust,请安装夜间版(rustup install nightly)。确保您拥有最新的夜间版(rustup update nightly)。最后,设置该项目使用夜间版(rustup override set nightly)。

您现在可以使用cargo run运行程序。该程序对32个大写字母的字节应用ROT13解密。借助SIMD,程序可以同时解密所有32个字节。

让我们看看程序的每个部分以了解其工作原理。它以以下方式开始:

#![feature(portable_simd)]use core::simd::prelude::*;

Rust夜间版仅在请求时提供其额外的功能(或“特性”)。#![feature(portable_simd)]语句请求Rust夜间版提供新的实验性core::simd模块。然后,use语句导入了该模块最重要的类型和特性。

在代码的下一个部分中,我们定义了有用的常量:

const LANES: usize = 32;const THIRTEENS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([13; LANES]);const TWENTYSIXS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([26; LANES]);const ZEES: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([b'Z'; LANES]);

Simd 结构是一种特殊的 Rust 数组。例如,它始终是内存对齐的。 常量 LANES 表示 Simd 数组的长度。 from_array 构造函数将常规的 Rust 数组复制到 Simd 中。 在这种情况下,因为我们想要 const Simd,我们构建的数组也必须是 const

接下来的两行将加密文本复制到 data 中,然后对每个字母加上 13。

let mut data = Simd::<u8, LANES>::from_slice(b"URYYBJBEYQVQBUBCRVGFNYYTBVATJRYY");data += THIRTEENS;

如果您犯了一个错误,导致加密文本的长度不是 LANES(32),那会怎么样?不幸的是,编译器不会告诉您。相反,当您运行程序时,from_slice 会引发错误。如果加密文本包含非大写字母呢?在这个示例程序中,我们将忽略这种可能性。

+= 运算符对 Simd dataSimd THIRTEENS 进行逐元素相加。它将结果放入 data 中。请记住,使用常规 Rust 加法的调试构建会检查溢出情况。但 SIMD 不会如此。 Rust 定义的 SIMD 算术运算符总是会进行环绕操作。类型为 u8 的值在 255 后会环绕。

巧合的是,Rot13 解密也需要进行环绕,但不是在 255 后,而是在 ‘Z’ 后。下面是一种实现所需的 Rot13 环绕的方法。对于超过 ‘Z’ 的任何值,它会将其减去 26。

let mask = data.simd_gt(ZEES);data = mask.select(data - TWENTYSIXS, data);

这表示寻找逐元素超过 ‘Z’ 的位置。然后,从所有值中减去 26。在感兴趣的位置,使用减去后的值。在其他位置,使用原始值。将所有值都减去然后仅使用一部分是否看起来很浪费?对于 SIMD 而言,这不需要额外的计算时间,并且避免了跳转。因此,这种策略是高效且常见的。

程序的结束部分如下:

let output = String::from_utf8_lossy(data.as_array());assert_eq!(output, "HELLOWORLDIDOHOPEITSALLGOINGWELL");println!("{}", output);

请注意 .as_array() 方法。它可以安全地将 Simd 结构转换为常规的 Rust 数组,而无需复制。

令我感到惊讶的是,即使在没有 SIMD 扩展的计算机上,该程序也能正常运行。Rust 夜版会将代码编译成常规(非 SIMD)指令。但我们不仅仅想要运行“正常”,我们还想要更快地运行。这需要我们开启计算机的 SIMD 功能。

法则 2:CCC:检查、控制和选择计算机的 SIMD 能力。

为了使 SIMD 程序在您的计算机上运行更快,您首先必须发现您的计算机支持哪些 SIMD 扩展。如果您有 Intel/AMD 计算机,可以使用我的 simd-detect cargo 命令。

运行:

rustup override set nightlycargo install cargo-simd-detect --forcecargo simd-detect

在我的计算机上,它输出:

extension       width                   available       enabledsse2            128-bit/16-bytes        true            trueavx2            256-bit/32-bytes        true            falseavx512f         512-bit/64-bytes        true            false

这说明我的计算机支持 sse2avx2avx512f SIMD 扩展。其中,默认情况下,Rust 启用了无处不在的二十年历史的 sse2 扩展。

SIMD扩展形成了一个层次结构,其中 avx512f 位于 avx2sse2 之上。 启用较高级别的扩展还会启用较低级别的扩展。

大多数 Intel/AMD 计算机也支持十年前的 avx2 扩展。 通过设置环境变量来启用它:

# 在 Windows 命令提示符下设置 RUSTFLAGS=-C target-feature=+avx2# 在类 Unix 的 shell(如 Bash)中设置export RUSTFLAGS="-C target-feature=+avx2"

“强制安装”并再次运行 simd-detect,您应该看到已启用 avx2

# 每次强制安装以查看 'enabled' 的更改cargo install cargo-simd-detect --forcecargo simd-detect

extension         width                   available       enabledsse2            128位/16字节        true            trueavx2            256位/32字节        true            trueavx512f         512位/64字节        true            false

或者,您可以打开机器支持的所有 SIMD 扩展:

# 在 Windows 命令提示符下设置 RUSTFLAGS=-C target-cpu=native# 在类 Unix 的 shell(如 Bash)中设置export RUSTFLAGS="-C target-cpu=native"

在我的机器上,这将启用 avx512f,一种由一些 Intel 计算机和少数 AMD 计算机支持的较新的 SIMD 扩展。

您可以使用以下方法将 SIMD 扩展恢复为默认值(在 Intel/AMD 上为 sse2):

# 在 Windows 命令提示符下设置 RUSTFLAGS=# 在类 Unix 的 shell(如 Bash)中取消设置 RUSTFLAGS

您可能想知道为什么 target-cpu=native不是 Rust 的默认值。问题在于使用 avx2avx512f 创建的二进制文件将无法在缺少这些 SIMD 扩展的计算机上运行。因此,如果您仅为自己使用进行编译,请使用 target-cpu=native。然而,如果您为他人进行编译,请谨慎选择您的 SIMD 扩展,并让人们知道您假定的 SIMD 扩展级别。

令人高兴的是,无论您选择哪个 SIMD 扩展级别,Rust 的 SIMD 支持都非常灵活,您可以轻松更改决策。接下来,让我们详细了解在 Rust 中使用 SIMD 进行编程的细节。

规则 3:选择性地学习 core::simd

要使用 Rust 的新 core::simd 模块构建,您应该学习所选的构建块。下面是一个包含我发现最有用的结构体、方法等的速查表。每个项目都包括一个指向文档的链接。

结构体

  • Simd – 一种特殊的、对齐的、固定长度的 SimdElement 数组。我们称数组中的位置和存储在该位置上的元素为“lane”。默认情况下,我们复制 Simd 结构而不是引用它们。
  • Mask – 一种特殊的布尔数组,按每个“lane”显示包含/排除。

SimdElements

  • 浮点类型: f32f64
  • 整数类型: i8u8i16u16i32u32i64u64isizeusize
  • 但不包括 i128 u128

Simd 构造函数

  • Simd::from_array – 通过复制固定长度的数组创建一个 Simd 结构。
  • Simd::from_slice – 通过复制切片的前 LANE 个元素创建一个 Simd<T,LANE> 结构。
  • Simd::splat – 在 Simd 结构的所有lanes中复制一个单一值。
  • slice::as_simd – 安全地将常规切片转换为 Simd 的对齐切片(以及不对齐部分),无需复制。

Simd 转换

  • Simd::as_array – 安全地将 Simd 结构转换为常规数组引用,无需复制。

Simd 方法和运算符

  • simd[i] – 从 Simd 的单个lane中提取一个值。
  • simd + simd – 对两个 Simd 结构进行逐元素加法。同样支持 -*/%、求余、位与、位或、异或、非、移位。
  • simd += simd – 在原地将另一个 Simd 结构添加到当前结构。其他运算符也支持。
  • Simd::simd_gt – 比较两个 Simd 结构,返回一个 Mask,指示第一个结构的哪些元素大于第二个结构的元素。也支持 simd_ltsimd_lesimd_gesimd_eqsimd_ne
  • Simd::rotate_elements_left – 将 Simd 结构的元素向左旋转指定量。也支持 rotate_elements_right
  • simd_swizzle!(simd, indexes) – 根据指定的常量索引重新排列 Simd 结构的元素。
  • simd == simd – 检查两个 Simd 结构是否相等,返回常规的 bool 结果。
  • Simd::reduce_and – 对 Simd 结构的所有lane进行位与归约操作。还支持: reduce_orreduce_xorreduce_maxreduce_minreduce_sum(但不支持reduce_eq)。

Mask 方法和运算符

  • Mask::select – 根据掩码从两个 Simd 结构中选择元素。
  • Mask::all – 判断掩码是否全部为 true
  • Mask::any – 判断掩码是否包含任何 true

关于lanes的一切

  • Simd::LANES – 一个常量,表示 Simd 结构中的元素(lanes)数量。
  • SupportedLaneCount – 指示 LANES 的允许值。在泛型中使用。
  • simd.lanes – 常量方法,用于获取 Simd 结构的lane数量。

低级别对齐、偏移等

当可能时,请使用to_simd

  • mem::size_ofmem::align_ofmem::align_tointrinsics::offsetpointer::read_unaligned(不安全),pointer::write_unaligned(不安全),mem::transmute(不安全,常量)

更多,可能感兴趣的内容

  • deinterleavegather_orreversescatter

有了这些基础知识,现在是时候构建一些东西了。

规则4:思考候选算法。

你想加快什么?你事先不会知道哪个SIMD方法(若有)可以提供最佳效果。因此,你应该创建许多算法,然后进行分析(规则5)和基准测试(规则7)。

我想加快range-set-blaze,这是一个用于操作“杂乱”的整数集的crate。我希望创建一个名为is_consecutive的函数来检测连续整数块,这可能很有用。

背景: Crate range-set-blaze 用于处理“杂乱”的整数。这里的“杂乱”是指用于表示数据所需的范围数量与输入整数的数量相比非常小。例如,这1002个输入整数

100, 101, …, 489, 499, 501, 502,…, 998, 999, 999, 100, 0

最终变为三个Rust范围:

0..=0, 100..=499, 501..=999

(在内部,RangeSetBlaze结构将一组整数表示为按顺序排列的不相交范围的缓存高效的BTreeMap。)

尽管输入整数允许无序和冗余,但我们希望它们通常是“好的”。RangeSetBlaze的from_iter构造函数已经利用了这一预期,通过分组相邻整数来处理。例如,from_iter首先将这1002个输入整数转换为四个范围

100..=499, 501..=999, 100..=100, 0..=0.

以最小的、常量的内存使用量进行,与输入大小无关。然后对这些缩小的范围进行排序和合并。

我想知道一个新的from_slice方法是否可以通过快速查找(一些)连续整数来加快从类似数组的输入构造的过程。例如,它是否可以使用最小的、常量的内存将这1002个输入整数转换为五个Rust范围:

100..=499, 501..=999, 999..=999, 100..=100, 0..=0.

如果可以的话,from_iter然后可以快速完成处理。

让我们先用常规的Rust编写is_consecutive

pub const LANES: usize = 16;pub fn is_consecutive_regular(chunk: &[u32; LANES]) -> bool {    for i in 1..LANES {        if chunk[i - 1].checked_add(1) != Some(chunk[i]) {            return false;        }    }    true}

该算法只是顺序地循环遍历数组,检查每个值是否比其前一个值大1。它还避免了溢出。

循环遍历这些项似乎很容易,我不确定SIMD是否能有所改善。这是我的第一次尝试:

Splat0

use std::simd::prelude::*;const COMPARISON_VALUE_SPLAT0: Simd<u32, LANES> =    Simd::from_array([15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]);pub fn is_consecutive_splat0(chunk: Simd<u32, LANES>) -> bool {    if chunk[0].overflowing_add(LANES as u32 - 1) != (chunk[LANES - 1], false) {        return false;    }    let added = chunk + COMPARISON_VALUE_SPLAT0;    Simd::splat(added[0]) == added}

这是其计算概述:

来源:此图片及后续所有图片来自作者。

首先(无用地)检查第一个和最后一个项目之间是否相差15。然后通过将15添加到第0个项目、14添加到下一个项目等来创建added。最后,为了查看added中的所有项是否相同,它基于added的第0个项创建了一个新的Simd,然后进行比较。请回忆一下splat从一个值创建一个Simd结构体。

Splat1和Splat2

当我向Ben Lichtman提到is_consecutive问题时,他独立提出了Splat1:

const COMPARISON_VALUE_SPLAT1: Simd<u32, LANES> =    Simd::from_array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]);pub fn is_consecutive_splat1(chunk: Simd<u32, LANES>) -> bool {    let subtracted = chunk - COMPARISON_VALUE_SPLAT1;    Simd::splat(chunk[0]) == subtracted}

Splat1从chunk中减去比较值,并检查结果是否与chunk的第一个元素(通过splat)相同。

为了加速你的Rust代码的SIMD,需要遵循九个规则(第一部分) 四海 第3张

他还提出了一个名为Splat2的变种,它将subtracted的第一个元素进行splat,而不是chunk。这似乎避免了一次内存访问。

我相信你会想知道其中哪种方法最好,但在我们讨论这个之前,让我们看看另外两个候选者。

Swizzle

Swizzle与Splat2类似,但使用simd_swizzle!代替splat。宏simd_swizzle!根据索引数组重新排列旧Simd的通道,创建一个新Simd

pub fn is_consecutive_sizzle(chunk: Simd<u32, LANES>) -> bool {    let subtracted = chunk - COMPARISON_VALUE_SPLAT1;    simd_swizzle!(subtracted, [0; LANES]) == subtracted}

Rotate

这个方法有所不同。我对它抱有很高的期望。

const COMPARISON_VALUE_ROTATE: Simd<u32, LANES> =    Simd::from_array([4294967281, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]);pub fn is_consecutive_rotate(chunk: Simd<u32, LANES>) -> bool {    let rotated = chunk.rotate_elements_right::<1>();    chunk - rotated == COMPARISON_VALUE_ROTATE}

这个思路是将所有元素向右旋转一次。然后我们从原始的chunk中减去rotated。如果输入是连续的,结果应该是“-15”,其后是所有1(使用环绕减法,-15是4294967281u32)。

为了加速你的Rust代码的SIMD,需要遵循九个规则(第一部分) 四海 第4张

现在我们有了候选人,让我们开始评估他们。

规则5:使用Godbolt和AI来理解你的代码的汇编,即使你不懂汇编语言。

我们将以两种方式评估候选人。首先,在这条规则中,我们将查看从我们的代码生成的汇编语言。其次,在规则7中,我们将对代码的速度进行基准测试。

如果你不懂汇编语言,不必担心,你仍然可以从中获得一些收获。

查看生成的汇编语言最简单的方法是使用Compiler Explorer(即Godbolt)。它对于不使用外部库的短代码效果最佳。它的界面如下所示:

为了加速你的Rust代码的SIMD,需要遵循九个规则(第一部分) 四海 第5张

参考上图中的数字,按照以下步骤使用Godbolt:

  1. 用您的浏览器打开godbolt.org
  2. 添加一个新的代码编辑器。
  3. 选择Rust作为您的语言。
  4. 粘贴感兴趣的代码。使感兴趣的函数公开可见(pub fn)。不要包括主函数或不需要的函数。该工具不支持外部库。
  5. 添加一个新的编译器。
  6. 设置编译器版本为nightly。
  7. 设置选项(暂时)为-C opt-level=3 -C target-feature=+avx512f
  8. 如果有错误,请查看输出。
  9. 如果您想共享或保存工具的状态,请点击“共享”

从上图中,您可以看到Splat2和Sizzle完全相同,所以我们可以将Sizzle排除在考虑范围之外。如果您打开我的Godbolt会话的副本,您还会看到大多数函数编译出大约相同数量的汇编操作。例外是Regular,它要长得多,以及包含早期检查的Splat0。

在汇编中,512位寄存器以ZMM开头。256位寄存器以YMM开头。128位寄存器以XMM开头。如果您想更好地理解生成的汇编语言,请使用AI工具生成注释。例如,这里我向Bing Chat询问了Splat2:

为了加速你的Rust代码的SIMD,需要遵循九个规则(第一部分) 四海 第6张

尝试不同的编译器设置,包括-C target-feature=+avx2,然后完全不设置target-feature

较少的汇编操作不一定意味着更快的速度。然而,查看汇编代码可以让我们明白编译器至少尝试使用SIMD操作、内联const引用等,并帮助我们判断两个候选人是否相同,就像Splat1和Swizzle一样。

您可能需要比Godbolt提供的反汇编功能更多的功能,例如能够处理使用外部库的代码。B3NNY向我推荐了cargo工具cargo-show-asm。我尝试了一下,发现使用起来相当容易。

range-set-blaze库必须处理超出u32的整型。此外,我们必须选择一定数量的LANES,但我们没有理由认为16 LANES总是最好的。为了满足这些需求,在下一个规则中,我们将推广代码。

规则6:通过内联泛型和无法使用时的宏以及无法使用时的特性,将通用化扩展到所有类型和LANE。

首先,用泛型来通用化Splat1。

#[inline]pub fn is_consecutive_splat1_gen<T, const N: usize>(    chunk: Simd<T, N>,    comparison_value: Simd<T, N>,) -> boolwhere    T: SimdElement + PartialEq,    Simd<T, N>: Sub<Simd<T, N>, Output = Simd<T, N>>,    LaneCount<N>: SupportedLaneCount,{    let subtracted = chunk - comparison_value;    Simd::splat(chunk[0]) == subtracted}

首先,注意到#[inline]属性。这对于效率很重要,我们将在这些小函数中几乎都使用它。

上面定义的函数is_consecutive_splat1_gen看起来很好,只是它需要一个叫做comparison_value的第二个输入,我们尚未定义。

如果你不需要一个通用的常量comparison_value,那我很羡慕你。如果你愿意,你可以跳到下一个规则。同样,如果你是从未来阅读这篇文章的,并且创建通用的常量comparison_value像你的个人机器人做家务一样简单,那我就更加羡慕你了。

我们可以尝试创建一个通用和常量的comparison_value_splat_gen。不幸的是,From<usize>和替代的T::One都不是常量,所以这行不通:

// 无法工作,因为From<usize>不是constpub const fn comparison_value_splat_gen<T, const N: usize>() -> Simd<T, N>where    T: SimdElement + Default + From<usize> + AddAssign,    LaneCount<N>: SupportedLaneCount,{    let mut arr: [T; N] = [T::from(0usize); N];    let mut i_usize = 0;    while i_usize < N {        arr[i_usize] = T::from(i_usize);        i_usize += 1;    }    Simd::from_array(arr)}

宏是流氓的最后避难所。所以,让我们使用宏:

#[macro_export]macro_rules! define_is_consecutive_splat1 {    ($function:ident, $type:ty) => {        #[inline]        pub fn $function<const N: usize>(chunk: Simd<$type, N>) -> bool        where            LaneCount<N>: SupportedLaneCount,        {            define_comparison_value_splat!(comparison_value_splat, $type);            let subtracted = chunk - comparison_value_splat();            Simd::splat(chunk[0]) == subtracted        }    };}#[macro_export]macro_rules! define_comparison_value_splat {    ($function:ident, $type:ty) => {        pub const fn $function<const N: usize>() -> Simd<$type, N>        where            LaneCount<N>: SupportedLaneCount,        {            let mut arr: [$type; N] = [0; N];            let mut i = 0;            while i < N {                arr[i] = i as $type;                i += 1;            }            Simd::from_array(arr)        }    };}

这使我们可以在任何特定的元素类型和所有LANES上运行(Rust Playground):

define_is_consecutive_splat1!(is_consecutive_splat1_i32, i32);let a: Simd<i32, 16> = black_box(Simd::from_array(array::from_fn(|i| 100 + i as i32)));let ninety_nines: Simd<i32, 16> = black_box(Simd::from_array([99; 16]));assert!(is_consecutive_splat1_i32(a));assert!(!is_consecutive_splat1_i32(ninety_nines));

可惜的是,这对于 range-set-blaze 来说仍然不够。它需要在所有元素类型(而不仅仅是一个)以及所有LANE(而不仅仅是一个)上运行。

幸运的是,有一个解决办法,再次依赖于宏。它还利用了一个事实,那就是我们只需要支持有限的类型列表,即:i8i16i32i64isizeu8u16u32u64usize。如果您还需要支持f32f64,那也没问题。

另一方面,如果您需要支持i128u128,您可能会没有办法。核心的simd模块不支持它们。我们将在第8条规则中看到range-set-blaze如何以性能成本解决这个问题。

这个解决方案定义了一个新的trait,这里称为IsConsecutive。然后我们使用一个宏(调用一个宏,再调用一个宏)在这10种感兴趣的类型上实现这个trait。

pub trait IsConsecutive {    fn is_consecutive<const N: usize>(chunk: Simd<Self, N>) -> bool    where        Self: SimdElement,        Simd<Self, N>: Sub<Simd<Self, N>, Output = Simd<Self, N>>,        LaneCount<N>: SupportedLaneCount;}macro_rules! impl_is_consecutive {    ($type:ty) => {        impl IsConsecutive for $type {            #[inline] // very important            fn is_consecutive<const N: usize>(chunk: Simd<Self, N>) -> bool            where                Self: SimdElement,                Simd<Self, N>: Sub<Simd<Self, N>, Output = Simd<Self, N>>,                LaneCount<N>: SupportedLaneCount,            {                define_is_consecutive_splat1!(is_consecutive_splat1, $type);                is_consecutive_splat1(chunk)            }        }    };}impl_is_consecutive!(i8);impl_is_consecutive!(i16);impl_is_consecutive!(i32);impl_is_consecutive!(i64);impl_is_consecutive!(isize);impl_is_consecutive!(u8);impl_is_consecutive!(u16);impl_is_consecutive!(u32);impl_is_consecutive!(u64);impl_is_consecutive!(usize);

现在我们可以调用完全通用的代码(Rust Playground):

// Works on i32 and 16 laneslet a: Simd<i32, 16> = black_box(Simd::from_array(array::from_fn(|i| 100 + i as i32)));let ninety_nines: Simd<i32, 16> = black_box(Simd::from_array([99; 16]));assert!(IsConsecutive::is_consecutive(a));assert!(!IsConsecutive::is_consecutive(ninety_nines));// Works on i8 and 64 laneslet a: Simd<i8, 64> = black_box(Simd::from_array(array::from_fn(|i| 10 + i as i8)));let ninety_nines: Simd<i8, 64> = black_box(Simd::from_array([99; 64]));assert!(IsConsecutive::is_consecutive(a));assert!(!IsConsecutive::is_consecutive(ninety_nines));

有了这种技术,我们可以创建多个完全通用的候选算法,它们可以处理类型和LANES。接下来,是时候进行基准测试,看哪些算法最快了。

这些是向Rust添加SIMD代码的前六条规则。在即将发布的第二部分中,我们将讨论第7到第9条规则。这些规则将涉及如何选择算法和设置LANES。还有如何将SIMD操作集成到您现有的代码中,以及如何使其可选。第2部分将在不久后发布。希望能在那里见到您。

关注VoAGI上的卡尔。我写有关Rust和Python科学编程、机器学习和统计的文章。我通常每月写一篇文章。

Leave a Reply

Your email address will not be published. Required fields are marked *