Take perf.

Author: jorgecarleitao

Cargo.toml

@@ -21,3 +21,7 @@ harness = false
 [[bench]]
 name = "sum_nulls_bitmap"
 harness = false
+
+[[bench]]
+name = "take"
+harness = false

benches/take.rs

@@ -0,0 +1,30 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+
+use simd_benches::take::*;
+
+fn close(l: &[f32], r: &[f32]) {
+    for (l, r) in l.iter().zip(r.iter()) {
+        assert!((l - r).abs() < l * 0.001);
+    }
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    let name = "";
+    (10..=20).step_by(2).for_each(|log2_size| {
+        let size = 2usize.pow(log2_size);
+        let array = (0..size).map(|x| 1.0 + x as f32).collect::<Vec<_>>();
+        let indices = (0..size).collect::<Vec<_>>();
+        let result = naive_take(&array, &indices);
+
+        c.bench_function(
+            &format!("core_simd_take{} 2^{} f32", name, log2_size),
+            |b| b.iter(|| close(&core_simd_take(&array, &indices), &result)),
+        );
+        c.bench_function(&format!("naive_take{} 2^{} f32", name, log2_size), |b| {
+            b.iter(|| close(&naive_take(&array, &indices), &result))
+        });
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

src/lib.rs

@@ -4,3 +4,4 @@ pub mod bitmap_ops;
 pub mod sum;
 pub mod sum_nulls;
 pub mod sum_nulls_bitmap;
+pub mod take;

src/take.rs

@@ -0,0 +1,25 @@
+use core_simd::*;
+
+pub fn naive_take(values: &[f32], indices: &[usize]) -> Vec<f32> {
+    indices.iter().map(|i| values[*i]).collect()
+}
+
+pub fn core_simd_take(values: &[f32], indices: &[usize]) -> Vec<f32> {
+    let chunks = indices.chunks_exact(8);
+    // todo handle remainder
+
+    let mut result = vec![0.0; indices.len()]; // todo: maybeUninit
+    let result_chunks = result.chunks_exact_mut(8);
+    chunks.zip(result_chunks).for_each(|(chunk, r_chunk)| {
+        let idxs: [usize; 8] = chunk.try_into().unwrap();
+        let idxs: usizex8 = usizex8::from_array(idxs);
+
+        let r = Simd::gather_or_default(&values, idxs);
+        let r: [f32; 8] = r.to_array();
+
+        let r_chunk: &mut [f32; 8] = r_chunk.try_into().unwrap();
+        *r_chunk = r;
+    });
+
+    result
+}