Switch to a more performance-oriented iterator implementation

Cargo.toml

@@ -9,9 +9,11 @@ edition = "2021"
 atomic-wait = { version = "1.1.0", default-features = false }
 crossbeam = { version = "0.8.2", features = ["crossbeam-channel", "crossbeam-deque", "crossbeam-epoch", "crossbeam-queue"] }
 hwlocality = { version = "1.0.0", path = "../hwlocality" }
+itertools = "0.11"
 
 [dev-dependencies]
 criterion = { version = "0.5.1", default-features = false }
+pessimize = { git = "https://github.com/HadrienG2/pessimize.git" }
 proptest = "1.3"
 rayon = { version = "1.8.0", default-features = false }
 

benches/benchmark.rs

@@ -1,45 +1,122 @@
 use criterion::{criterion_group, criterion_main, Bencher, BenchmarkId, Criterion};
-use sched_local::FlatPool;
-use std::hint::black_box;
+use sched_local::{flags::AtomicFlags, FlatPool};
+use std::sync::atomic::Ordering;
 
-fn fibonacci_rayon(n: u64) -> u64 {
-    if n > 1 {
-        let (x, y) = rayon::join(|| fibonacci_rayon(n - 1), || fibonacci_rayon(n - 2));
-        x + y
-    } else {
-        n
+fn criterion_benchmark(c: &mut Criterion) {
+    bench_flags(c);
+    bench_fibonacci(c);
+}
+
+fn bench_flags(c: &mut Criterion) {
+    for len_pow2 in 0..=8 {
+        let len = 2usize.pow(len_pow2);
+        let flags = AtomicFlags::new(len);
+        let header = format!("flags/{len}");
+
+        // General logic for benchmarks that target a specific index
+        fn bench_indexed_op<R>(
+            c: &mut Criterion,
+            flags: &AtomicFlags,
+            group_name: &str,
+            mut op: impl FnMut(&AtomicFlags, usize) -> R,
+        ) {
+            use pessimize::hide;
+            let mut group = c.benchmark_group(group_name);
+            group.bench_function("first", |b| b.iter(|| op(hide(flags), hide(0))));
+            group.bench_function("center", |b| {
+                let center = flags.len() / 2;
+                b.iter(|| op(hide(flags), hide(center)))
+            });
+            group.bench_function("last", |b| {
+                let last = flags.len() - 1;
+                b.iter(|| op(hide(flags), hide(last)))
+            });
+        }
+
+        // Operations that test a single bit
+        bench_indexed_op(c, &flags, &format!("{header}/is_set"), |flags, pos| {
+            flags.is_set(pos, Ordering::Relaxed)
+        });
+        bench_indexed_op(c, &flags, &format!("{header}/fetch_set"), |flags, pos| {
+            flags.fetch_set(pos, Ordering::Relaxed)
+        });
+        bench_indexed_op(c, &flags, &format!("{header}/fetch_clear"), |flags, pos| {
+            flags.fetch_clear(pos, Ordering::Relaxed)
+        });
+
+        // Operations that set all bits to the same value
+        c.bench_function(&format!("{header}/set_all"), |b| {
+            b.iter(|| pessimize::hide(&flags).set_all(Ordering::Relaxed))
+        });
+        c.bench_function(&format!("{header}/clear_all"), |b| {
+            b.iter(|| pessimize::hide(&flags).clear_all(Ordering::Relaxed))
+        });
+
+        // Run one iteration of the iterators over set and unset indices
+        bench_indexed_op(
+            c,
+            &flags,
+            &format!("{header}/iter_set_around/once"),
+            |flags, pos| flags.iter_set_around(pos, Ordering::Relaxed).next(),
+        );
+        bench_indexed_op(
+            c,
+            &flags,
+            &format!("{header}/iter_set_around/all"),
+            |flags, pos| flags.iter_set_around(pos, Ordering::Relaxed).count(),
+        );
+        bench_indexed_op(
+            c,
+            &flags,
+            &format!("{header}/iter_unset_around/once"),
+            |flags, pos| flags.iter_unset_around(pos, Ordering::Relaxed).next(),
+        );
+        bench_indexed_op(
+            c,
+            &flags,
+            &format!("{header}/iter_unset_around/all"),
+            |flags, pos| flags.iter_unset_around(pos, Ordering::Relaxed).count(),
+        );
     }
 }
 
-pub fn criterion_benchmark(c: &mut Criterion) {
+fn bench_fibonacci(c: &mut Criterion) {
     fn bench_backend(
         c: &mut Criterion,
         backend_name: &str,
-        mut fibonacci_bench: impl FnMut(&mut Bencher, u64),
+        mut bench_impl: impl FnMut(&mut Bencher, u64),
     ) {
-        let group_name = format!("{backend_name}/fibonacci");
+        let group_name = format!("fibonacci/{backend_name}");
         let mut group = c.benchmark_group(group_name);
         for size in [1, 2, 4, 8, 16, 20, 24, 28, 30, 32, 34] {
             let phi = (1.0 + 5.0f64.sqrt()) / 2.0;
             group.throughput(criterion::Throughput::Elements(
                 phi.powi(i32::try_from(size).unwrap()) as u64,
             ));
             group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, size| {
-                fibonacci_bench(b, *size)
+                bench_impl(b, *size)
             });
         }
     }
 
     bench_backend(c, "rayon", |b: &mut Bencher, size| {
-        b.iter(|| fibonacci_rayon(black_box(size)))
+        b.iter(|| fibonacci_rayon(pessimize::hide(size)))
     });
 
     let pool = FlatPool::new();
-    // FIXME: Try to use bench_with_input so I don't need to recreate the scope
     bench_backend(c, "flat", |b: &mut Bencher, size| {
-        pool.scope(|scope| b.iter(|| sched_local::fibonacci_flat(scope, size)))
+        pool.scope(|scope| b.iter(|| sched_local::fibonacci_flat(scope, pessimize::hide(size))))
     })
 }
 
+fn fibonacci_rayon(n: u64) -> u64 {
+    if n > 1 {
+        let (x, y) = rayon::join(|| fibonacci_rayon(n - 1), || fibonacci_rayon(n - 2));
+        x + y
+    } else {
+        n
+    }
+}
+
 criterion_group!(benches, criterion_benchmark);
 criterion_main!(benches);