Nine Rules for SIMD Acceleration of Your Rust Code (Part 1)
\n \t
\n
<\/BR>
\n
\n \t
\n
<\/BR><\/p>\n
\nThanks to Ben Lichtman (B3NNY) at the Seattle Rust Meetup for pointing me in the right direction on SIMD.<\/em><\/p>\n<\/blockquote>\n
SIMD<\/a>\u00a0(Single Instruction, Multiple Data) operations have been a feature of Intel\/AMD and ARM CPUs since the early 2000s. These operations enable you to, for example, add an array of eight\u00a0
i32<\/code>\u00a0to another array of eight\u00a0i32<\/code>\u00a0with just one CPU operation\u00a0on a single core<\/strong>. Using SIMD operations greatly speeds up certain tasks. If you\u2019re not using SIMD, you may not be fully using your CPU\u2019s capabilities.<\/p>\nIs this \u201cYet Another Rust<\/a> and SIMD\u201d article? Yes and no. Yes, I did apply SIMD to a programming problem and then feel compelled to write an article about it. No, I hope that this article also goes into enough depth that it can guide you through\u00a0your<\/em>\u00a0project. It explains the newly available SIMD capabilities and settings in Rust nightly. It includes a Rust SIMD cheatsheet. It shows how to make your SIMD code generic without leaving safe Rust. It gets you started with tools such as Godbolt and Criterion. Finally, it introduces new cargo commands that make the process easier.<\/p>\n
\nThe\u00a0
range-set-blaze<\/code><\/a>\u00a0crate uses its\u00a0RangeSetBlaze::from_iter<\/code>\u00a0method to ingest potentially long sequences of integers. When the integers are \u201cclumpy\u201d, it can do this\u00a030 times faster<\/a>\u00a0than Rust\u2019s standard\u00a0HashSet::from_iter<\/code>. Can we do even better if we use Simd<\/a> operations? Yes!<\/p>\n\nSee\u00a0this documentation<\/a>\u00a0for the definition of \u201cclumpy\u201d. Also, what happens if the integers are not clumpy?\u00a0
RangeSetBlaze<\/code>\u00a0is\u00a02 to 3 times\u00a0<\/a><\/em>slower<\/a>\u00a0than\u00a0HashSet<\/code>.<\/em><\/p>\n<\/blockquote>\nOn clumpy integers,\u00a0
RangeSetBlaze::from_slice\u00a0<\/code>\u2014 a new method based on SIMD operations \u2014 is 7 times faster than\u00a0RangeSetBlaze::from_iter.\u00a0<\/code>That makes it more than 200 times faster than\u00a0HashSet::from_iter<\/code>. (When the integers are not clumpy, it is still 2 to 3 times slower than\u00a0HashSet<\/code>.)<\/p>\nOver the course of implementing this speed up, I learned nine rules that can help you accelerate your projects with SIMD operations.<\/p>\n
The rules are:<\/p>\n
\n
- Use nightly Rust and\u00a0
core::simd<\/code>, Rust\u2019s experimental standard SIMD module.<\/li>\n- CCC: Check, Control, and Choose your computer\u2019s SIMD capabilities.<\/li>\n
- Learn\u00a0
core::simd<\/code>, but selectively.<\/li>\n- Brainstorm candidate algorithms.<\/li>\n
- Use Godbolt and AI to understand your code\u2019s assembly, even if you don\u2019t know assembly language.<\/li>\n
- Generalize to all types and LANES with in-lined generics, (and when that doesn\u2019t work) macros, and (when that doesn\u2019t work) traits.<\/li>\n<\/ol>\n
See\u00a0Part 2<\/a>\u00a0for these rules:<\/p>\n
7. Use Criterion benchmarking to pick an algorithm and to discover that LANES should (almost) always be 32 or 64.<\/em><\/p>\n
8. Integrate your best SIMD algorithm into your project with\u00a0
as_simd<\/code>, special code for\u00a0i128\/u128<\/code>, and additional in-context benchmarking.<\/em><\/p>\n9. Extricate your best SIMD algorithm from your project (for now) with an optional cargo feature.<\/em><\/p>\n
Aside: To avoid wishy-washiness, I call these \u201crules\u201d, but they are, of course, just suggestions.<\/em><\/p>\n
Rule 1: Use nightly Rust and\u00a0
core::simd<\/code>, Rust\u2019s experimental standard SIMD module.<\/h2>\nRust can access SIMD operations either via the stable\u00a0
core::arch<\/code><\/a>\u00a0module or via nighty\u2019s\u00a0core::simd<\/code><\/a>\u00a0module. Let\u2019s compare them:<\/p>\n
core::arch<\/code><\/strong><\/p>\n\n
- Stable<\/li>\n
- \n\u201c[N]ot the easiest thing in the world<\/a>\u201d<\/li>\n
- Offers high-performance to downstream users of your crate. For example, because\u00a0regex<\/a>\u00a0and\u00a0
memchr<\/code><\/a>\u00a0went this route, over 100,000 other crates got stable SIMD acceleration for free. [Reddit discussion<\/a>,\u00a0some relevant\u00a0memchr<\/code>\u00a0code<\/a>]<\/li>\n<\/ul>\n
core::simd<\/code><\/strong><\/p>\n\n
- Nightly<\/li>\n
- Delightfully easy and portable.<\/li>\n
- Limits downstream users to nightly.<\/li>\n<\/ul>\n
I decided to go with \u201ceasy\u201d. If you decide to take the harder road, starting first with the easier path may still be worthwhile.<\/p>\n
\nIn either case, before we try to use SIMD operations in a larger project, let\u2019s make sure we can get them working at all. Here are the steps:<\/p>\n
First, create a project called\u00a0
simd_hello<\/code>:<\/p>\ncargo new simd_hello\ncd simd_hello<\/code><\/pre>\nEdit\u00a0
src\/main.rs<\/code>\u00a0to contain (Rust playground<\/a>):<\/p>\n\/\/ Tell nightly Rust to enable 'portable_simd'\n#![feature(portable_simd)]\nuse core::simd::prelude::*;\n\n\/\/ constant Simd structs\nconst LANES: usize = 32;\nconst THIRTEENS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([13; LANES]);\nconst TWENTYSIXS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([26; LANES]);\nconst ZEES: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([b'Z'; LANES]);\n\nfn main() {\n \/\/ create a Simd struct from a slice of LANES bytes\n let mut data = Simd::<u8, LANES>::from_slice(b\"URYYBJBEYQVQBUBCRVGFNYYTBVATJRYY\");\n\n data += THIRTEENS; \/\/ add 13 to each byte\n\n \/\/ compare each byte to 'Z', where the byte is greater than 'Z', subtract 26\n let mask = data.simd_gt(ZEES); \/\/ compare each byte to 'Z'\n data = mask.select(data - TWENTYSIXS, data);\n\n let output = String::from_utf8_lossy(data.as_array());\n assert_eq!(output, \"HELLOWORLDIDOHOPEITSALLGOINGWELL\");\n println!(\"{}\", output);\n}<\/code><\/pre>\nNext \u2014 full SIMD capabilities require the nightly version of Rust. Assuming you have Rust installed, install nightly (
rustup install nightly<\/code>). Make sure you have the latest nightly version (rustup update nightly<\/code>). Finally, set this project to use nightly (rustup override set nightly<\/code>).<\/p>\nYou can now run the program with\u00a0
cargo run<\/code>. The program applies\u00a0ROT13 decryption<\/a>\u00a0to 32 bytes of upper-case letters. With SIMD, the program can decrypt all 32 bytes simultaneously.<\/p>\nLet\u2019s look at each section of the program to see how it works. It starts with:<\/p>\n
#![feature(portable_simd)]\nuse core::simd::prelude::*;<\/code><\/pre>\nRust nightly offers its extra capabilities (or \u201cfeatures\u201d) only on request. The\u00a0
#![feature(portable_simd)]<\/code>\u00a0statement requests that Rust nightly make available the new experimental\u00a0core::simd<\/code>\u00a0module. The\u00a0use<\/code>\u00a0statement then imports the module\u2019s most important types and traits.<\/p>\nIn the code\u2019s next section, we define useful constants:<\/p>\n
const LANES: usize = 32;\nconst THIRTEENS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([13; LANES]);\nconst TWENTYSIXS: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([26; LANES]);\nconst ZEES: Simd<u8, LANES> = Simd::<u8, LANES>::from_array([b'Z'; LANES]);<\/code><\/pre>\nThe\u00a0
Simd<\/code>\u00a0struct is a special kind of Rust array. (It is, for example, always memory aligned.) The constant\u00a0LANES<\/code>\u00a0tells the length of the\u00a0Simd<\/code>\u00a0array. The\u00a0from_array<\/code>\u00a0constructor copies a regular Rust array to create a\u00a0Simd<\/code>. In this case, because we want\u00a0const<\/code>\u00a0Simd<\/code>\u2019s, the arrays we construct from must also be\u00a0const<\/code>.<\/p>\nThe next two lines copy our encrypted text into\u00a0
data<\/code>\u00a0and then adds 13 to each letter.<\/p>\nlet mut data = Simd::<u8, LANES>::from_slice(b\"URYYBJBEYQVQBUBCRVGFNYYTBVATJRYY\");\ndata += THIRTEENS;<\/code><\/pre>\nWhat if you make an error and your encrypted text isn\u2019t exactly length\u00a0
LANES<\/code>\u00a0(32)? Sadly, the compiler won\u2019t tell you. Instead, when you run the program,\u00a0from_slice<\/code>\u00a0will panic. What if the encrypted text contains non-upper-case letters? In this example program, we\u2019ll ignore that possibility.<\/p>\nThe\u00a0
+=<\/code>\u00a0operator does element-wise addition between the\u00a0Simd<\/code>\u00a0data<\/code>\u00a0and\u00a0Simd<\/code>\u00a0THIRTEENS<\/code>. It puts the result in\u00a0data<\/code>. Recall that debug builds of regular Rust addition check for overflows. Not so with SIMD. Rust defines SIMD arithmetic operators to always wrap. Values of type\u00a0u8<\/code>\u00a0wrap after 255.<\/p>\nCoincidentally, Rot13 decryption also requires wrapping, but after \u2018Z\u2019 rather than after 255. Here is one approach to coding the needed Rot13 wrapping. It subtracts 26 from any values\u00a0on beyond \u2018Z<\/a>\u2019.<\/p>\n
let mask = data.simd_gt(ZEES);\ndata = mask.select(data - TWENTYSIXS, data);<\/code><\/pre>\nThis says to find the element-wise places beyond \u2018Z\u2019. Then, subtract 26 from all values. At the places of interest, use the subtracted values. At the other places, use the original values. Does subtracting from all values and then using only some seem wasteful? With SIMD, this takes no extra computer time and avoids jumps. This strategy is, thus, efficient and common.<\/p>\n
The program ends like so:<\/p>\n
let output = String::from_utf8_lossy(data.as_array());\nassert_eq!(output, \"HELLOWORLDIDOHOPEITSALLGOINGWELL\");\nprintln!(\"{}\", output);<\/code><\/pre>\nNotice the\u00a0
.as_array()<\/code>\u00a0method. It safely transmutes a\u00a0Simd<\/code>\u00a0struct into a regular Rust array without copying.<\/p>\nSurprisingly to me, this program runs fine on computers without SIMD extensions. Rust nightly compiles the code to regular (non-SIMD) instructions. But we don\u2019t just want to run \u201cfine\u201d, we want to run\u00a0faster<\/em>. That requires us to turn on our computer\u2019s SIMD power.<\/p>\n
Rule 2: CCC: Check, Control, and Choose your computer\u2019s SIMD capabilities.<\/h2>\n
To make SIMD programs run faster on your machine, you must first discover which SIMD extensions your machine supports. If you have an Intel\/AMD machine, you can use my\u00a0
simd-detect<\/code><\/a>\u00a0cargo command.<\/p>\nRun with:<\/p>\n
rustup override set nightly\ncargo install cargo-simd-detect --force\ncargo simd-detect<\/code><\/pre>\nOn my machine, it outputs:<\/p>\n
extension width available enabled\nsse2 128-bit\/16-bytes true true\navx2 256-bit\/32-bytes true false\navx512f 512-bit\/64-bytes true false<\/code><\/pre>\nThis says that my machine supports the\u00a0
sse2<\/code>,\u00a0avx2<\/code>, and\u00a0avx512f<\/code>\u00a0SIMD extensions. Of those, by default, Rust enables the ubiquitous twenty-year-old\u00a0sse2<\/code>\u00a0extension.<\/p>\nThe SIMD extensions form a hierarchy with\u00a0
avx512f<\/code>\u00a0above\u00a0avx2<\/code>\u00a0above\u00a0sse2<\/code>. Enabling a higher-level extension also enables the lower-level extensions.<\/p>\nMost Intel\/AMD computers also support the ten-year-old\u00a0
avx2<\/code>\u00a0extension. You enable it by setting an environment variable:<\/p>\n# For Windows Command Prompt\nset RUSTFLAGS=-C target-feature=+avx2\n\n# For Unix-like shells (like Bash)\nexport RUSTFLAGS=\"-C target-feature=+avx2\"<\/code><\/pre>\n\u201cForce install\u201d and run\u00a0
simd-detect<\/code>\u00a0again and you should see that\u00a0avx2<\/code>\u00a0is enabled.<\/p>\n# Force install every time to see changes to 'enabled'\ncargo install cargo-simd-detect --force\ncargo simd-detect<\/code><\/pre>\nextension width available enabled\nsse2 128-bit\/16-bytes true true\navx2 256-bit\/32-bytes true true\navx512f 512-bit\/64-bytes true false<\/code><\/pre>\nAlternatively, you can turn on every SIMD extension that your machine supports:<\/p>\n
# For Windows Command Prompt\nset RUSTFLAGS=-C target-cpu=native\n\n# For Unix-like shells (like Bash)\nexport RUSTFLAGS=\"-C target-cpu=native\"<\/code><\/pre>\nOn my machine this enables\u00a0
avx512f<\/code>, a newer SIMD extension supported by some Intel computers and a few AMD computers.<\/p>\nYou can set SIMD extensions back to their default (
sse2<\/code>\u00a0on Intel\/AMD) with:<\/p>\n# For Windows Command Prompt\nset RUSTFLAGS=\n\n# For Unix-like shells (like Bash)\nunset RUSTFLAGS<\/code><\/pre>\nYou may wonder why\u00a0
target-cpu=native<\/code>\u00a0isn\u2019t Rust\u2019s default. The problem is that binaries created using\u00a0avx2<\/code>\u00a0or\u00a0avx512f<\/code>\u00a0won\u2019t run on computers missing those SIMD extensions. So, if you are compiling only for your own use, use\u00a0target-cpu=native<\/code>. If, however, you are compiling for others, choose your SIMD extensions thoughtfully and let people know which SIMD extension level you are assuming.<\/p>\nHappily, whatever level of SIMD extension you pick, Rust\u2019s SIMD support is so flexible you can easily change your decision later. Let\u2019s next learn details of programming with SIMD in Rust.<\/p>\n
Rule 3: Learn\u00a0
core::simd<\/code>, but selectively.<\/h2>\nTo build with Rust\u2019s new\u00a0
core::simd<\/code><\/a>\u00a0module you should learn selected building blocks. Here is a\u00a0cheatsheet<\/a>\u00a0with the structs, methods, etc., that I\u2019ve found most useful. Each item includes a link to its\u00a0documentation<\/a>.<\/p>\nStructs<\/h3>\n
\n
- \n
Simd<\/code><\/a>\u00a0\u2013 a special, aligned, fixed-length array of\u00a0SimdElement<\/code><\/a>. We refer to a position in the array and the element stored at that position as a \u201clane\u201d. By default, we copy\u00a0Simd<\/code>\u00a0structs rather than reference them.<\/li>\n- \n
Mask<\/code><\/a>\u00a0\u2013 a special Boolean array showing inclusion\/exclusion on a per-lane basis.<\/li>\n<\/ul>\nSimdElements<\/h3>\n
\n
- Floating-Point Types:\u00a0
f32<\/code>,\u00a0f64<\/code>\n<\/li>\n- Integer Types:\u00a0
i8<\/code>,\u00a0u8<\/code>,\u00a0i16<\/code>,\u00a0u16<\/code>,\u00a0i32<\/code>,\u00a0u32<\/code>,\u00a0i64<\/code>,\u00a0u64<\/code>,\u00a0isize<\/code>,\u00a0usize<\/code>\n<\/li>\n- \u2014\u00a0but not\u00a0
i128<\/code>,\u00a0u128<\/code><\/em><\/a>\n<\/li>\n<\/ul>\n
Simd<\/code>\u00a0constructors<\/strong><\/h3>\n\n
- \n
Simd::from_array<\/code><\/a>\u00a0\u2013 creates a\u00a0Simd<\/code>\u00a0struct by copying a fixed-length array.<\/li>\n- \n
Simd::from_slice<\/code><\/a>\u00a0\u2013 creates a\u00a0Simd<T,LANE><\/code>\u00a0struct by copying the first\u00a0LANE<\/code>\u00a0elements of a slice.<\/li>\n- \n
Simd::splat<\/code><\/a>\u00a0\u2013 replicates a single value across all lanes of a\u00a0Simd<\/code>\u00a0struct.<\/li>\n- \n
slice::as_simd<\/code><\/a>\u00a0\u2013 without copying, safely transmutes a regular slice into an aligned slice of\u00a0Simd<\/code>\u00a0(plus unaligned leftovers).<\/li>\n<\/ul>\n
Simd<\/code>\u00a0conversion<\/strong><\/h3>\n