Created
March 25, 2023 10:15
-
-
Save dnutiu/8fc529d09d2872595cf2bf23d40e1cd2 to your computer and use it in GitHub Desktop.
Rust SIMD example; Requires nightly build and rand = "0.8.5" dependency.
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(test)] | |
use std::simd::{i32x16, Simd, SimdPartialEq, ToBitMask}; | |
use rand::Rng; | |
#[inline(never)] | |
fn hamming_distance_simd(a: &mut Vec<i32>, b: &mut Vec<i32>) -> u32 { | |
let mut distance: u32 = 0; | |
let lanes = 16; | |
let chunks = a.len() / lanes; | |
for i in 0..chunks { | |
// create simd vectors | |
let a: i32x16 = Simd::from_slice(&a[i*lanes..i*lanes+lanes]); | |
let b: i32x16 = Simd::from_slice(&b[i*lanes..i*lanes+lanes]); | |
// simd XOR | |
let z = a.simd_ne(b); | |
distance += z.to_bitmask().count_ones(); | |
} | |
return distance; | |
} | |
#[inline(never)] | |
fn hamming_distance_classic(a: &mut Vec<i32>, b: &mut Vec<i32>) -> u64 { | |
let mut distance = 0; | |
let length = a.len(); | |
for i in 0..length { | |
if a[i] != b[i] { | |
distance += 1 | |
} | |
} | |
return distance; | |
} | |
fn main() { | |
let mut a = Vec::new(); | |
let mut b = Vec::new(); | |
let mut rng = rand::thread_rng(); | |
for _ in 0..128 { | |
a.push(rng.gen_range(0..10)); | |
b.push(rng.gen_range(0..10)); | |
} | |
println!("SIMD"); | |
let distance = hamming_distance_simd(&mut a, &mut b); | |
println!("Distance {distance}"); | |
println!("CLASSIC"); | |
let distance = hamming_distance_classic(&mut a, &mut b); | |
println!("Distance {distance}"); | |
} | |
mod bench { | |
extern crate test; | |
use self::test::Bencher; | |
use std::iter; | |
static BENCH_SIZE: usize = 128; | |
macro_rules! bench { | |
($name:ident, $func:ident) => { | |
#[bench] | |
fn $name(b: &mut Bencher) { | |
let mut x: Vec<_> = iter::repeat(1i32) | |
.take(BENCH_SIZE) | |
.collect(); | |
let mut y: Vec<_> = iter::repeat(1i32) | |
.take(BENCH_SIZE) | |
.collect(); | |
b.iter(|| { | |
super::$func(&mut x, &mut y); | |
}) | |
} | |
} | |
} | |
bench!(simd, hamming_distance_simd); | |
bench!(vanilla, hamming_distance_classic); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment