Skip to content

Instantly share code, notes, and snippets.

@SoftPoison
Created June 10, 2024 09:14
Show Gist options
  • Save SoftPoison/07e10af46d17a7e1cf3cba9d4b23c0b8 to your computer and use it in GitHub Desktop.
Save SoftPoison/07e10af46d17a7e1cf3cba9d4b23c0b8 to your computer and use it in GitHub Desktop.
Fun little poking round with Rust SIMD and mmapping to make a silly little wordcount clone
#![feature(portable_simd)]
#![feature(slice_as_chunks)]
use std::{
ops::Neg,
simd::{cmp::SimdPartialEq, num::SimdInt},
};
use memmap2::Mmap;
const LANE_SIZE: usize = 64; // AVX-512 = 64 lanes
type ParallelU8 = std::simd::Simd<u8, LANE_SIZE>;
fn main() {
let filename = std::env::args().nth(1).expect("Usage: wctest <file>");
let file = std::fs::File::open(filename).expect("Failed to open file");
let mmap = unsafe { Mmap::map(&file) }.expect("Failed to mmap file");
let mask = ParallelU8::splat(0x0a);
// marginally faster unsafe version:
// let offset = mmap.as_ptr();
// let mut count = 0;
// for i in (0..mmap.len()).step_by(LANE_SIZE) {
// let x = unsafe { *offset.add(i).cast::<ParallelU8>() };
// count += x.simd_eq(mask).to_int().reduce_sum().neg() as usize;
// }
// usually there'll be enough zero bytes after the mapped region, buuuuuut this is just being hopeful
// good enough version that's doesn't use `unsafe`
let (chunks, rest) = mmap.as_chunks::<LANE_SIZE>();
let mut count = chunks
.iter()
.map(|x| {
ParallelU8::from_array(*x)
.simd_eq(mask)
.to_int()
.reduce_sum()
.neg() as usize
})
.reduce(|a, b| a + b)
.unwrap_or(0);
for b in rest {
if *b == 0x0a {
count += 1;
}
}
println!("{count}");
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment