Created
June 10, 2024 09:14
-
-
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
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#![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