feat(benches): implement thread binding with hwlocality (#337)

* add hwlocality dep & gating feature

* add utils function for thread binding

* add rayon threadpool init

* add std threads binding in `cut-edges`

* fix clippy lint

* remove feat from default

* disable --all-features in CI

Author: imrn99

.github/workflows/build.yml

@@ -23,7 +23,7 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Build crates
-        run: cargo build --all --all-features
+        run: cargo build --all # --all-features
   build_examples:
     name: Build examples for ${{ matrix.os }}
     strategy:
@@ -51,4 +51,4 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Build benchmarks
-        run: cargo build --benches --all-features
+        run: cargo build --benches # --all-features

.github/workflows/rust-test.yml

@@ -29,7 +29,7 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Run Clippy
-        run: cargo clippy --all-features -- -D warnings
+        run: cargo clippy # --all-features -- -D warnings
   tests:
     runs-on: ubuntu-22.04
     steps:
@@ -39,4 +39,4 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Run Tests
-        run: cargo test --all --all-features
+        run: cargo test --all # --all-features

Cargo.toml

@@ -54,6 +54,7 @@ smallvec = "2.0.0-alpha.10"
 
 # benchmarks
 criterion = "0.6.0"
+hwlocality = "1.0.0-alpha.7"
 iai-callgrind = "0.14.0"
 rand = "0.9.0-alpha.2"
 

benches/Cargo.toml

@@ -10,15 +10,18 @@ authors.workspace = true
 publish = false
 
 [features]
+# default = ["bind-threads"] # enable back when CI is fixed with all the deps
 _single_precision = []
+bind-threads = ["dep:hwlocality"]
 profiling = []
 
 # deps
 
 [dependencies]
-clap = { workspace =true, features = ["derive"] }
+clap = { workspace = true, features = ["derive"] }
 cfg-if.workspace = true
 honeycomb.workspace = true
+hwlocality = { workspace = true, optional = true }
 rayon.workspace = true
 rand = { workspace = true, features = ["small_rng"] }
 

benches/src/cut_edges.rs

@@ -11,17 +11,25 @@ use honeycomb::{
     prelude::{CMap2, CMapBuilder, CoordsFloat, DartIdType, EdgeIdType},
 };
 
-use crate::{cli::CutEdgesArgs, prof_start, prof_stop, utils::hash_file};
+use crate::{
+    cli::CutEdgesArgs,
+    prof_start, prof_stop,
+    utils::{get_num_threads, hash_file},
+};
 
 // const MAX_RETRY: u8 = 10;
 
 pub fn bench_cut_edges<T: CoordsFloat>(args: CutEdgesArgs) -> CMap2<T> {
     let input_map = args.input.to_str().unwrap();
     let target_len = T::from(args.target_length).unwrap();
 
-    let n_threads = std::thread::available_parallelism()
-        .map(|n| n.get())
-        .unwrap_or(1);
+    let n_threads = if let Ok(val) = get_num_threads() {
+        val
+    } else {
+        std::thread::available_parallelism()
+            .map(|n| n.get())
+            .unwrap_or(1)
+    };
 
     // load map from file
     let mut instant = Instant::now();
@@ -215,25 +223,77 @@ fn dispatch_std_threads<T: CoordsFloat>(
         .zip(darts.chunks(6))
         .map(|(e, sl)| (e, sl.try_into().unwrap()))
         .collect();
-    std::thread::scope(|s| {
-        let mut handles = Vec::new();
-        for wl in units.chunks(1 + units.len() / n_threads) {
-            handles.push(s.spawn(|| {
-                let mut n = 0;
-                wl.iter().for_each(|&(e, new_darts)| {
-                    let mut n_retry = 0;
-                    if map.is_i_free::<2>(e as DartIdType) {
-                        while !process_outer_edge(map, &mut n_retry, e, new_darts).is_validated() {}
-                    } else {
-                        while !process_inner_edge(map, &mut n_retry, e, new_darts).is_validated() {}
+
+    #[cfg(feature = "bind-threads")]
+    {
+        use std::sync::Arc;
+
+        use hwlocality::{Topology, cpu::binding::CpuBindingFlags};
+
+        use crate::utils::get_proc_list;
+
+        let topo = Arc::new(Topology::new().unwrap());
+        let mut cores = get_proc_list(&topo).unwrap_or_default().into_iter().cycle();
+        std::thread::scope(|s| {
+            let mut handles = Vec::new();
+            for wl in units.chunks(1 + units.len() / n_threads) {
+                let topo = topo.clone();
+                let core = cores.next();
+                handles.push(s.spawn(move || {
+                    // bind
+                    if let Some(c) = core {
+                        let tid = hwlocality::current_thread_id();
+                        topo.bind_thread_cpu(tid, &c, CpuBindingFlags::empty())
+                            .unwrap();
                     }
-                    n += n_retry as u32;
-                });
-                n
-            })); // s.spawn
-        } // for wl in workloads
-        handles.into_iter().map(|h| h.join().unwrap()).sum()
-    }) // std::thread::scope
+                    // work
+                    let mut n = 0;
+                    wl.iter().for_each(|&(e, new_darts)| {
+                        let mut n_retry = 0;
+                        if map.is_i_free::<2>(e as DartIdType) {
+                            while !process_outer_edge(map, &mut n_retry, e, new_darts)
+                                .is_validated()
+                            {}
+                        } else {
+                            while !process_inner_edge(map, &mut n_retry, e, new_darts)
+                                .is_validated()
+                            {}
+                        }
+                        n += n_retry as u32;
+                    });
+                    n
+                })); // s.spawn
+            } // for wl in workloads
+            handles.into_iter().map(|h| h.join().unwrap()).sum()
+        }) // std::thread::scope
+    }
+
+    #[cfg(not(feature = "bind-threads"))]
+    {
+        std::thread::scope(|s| {
+            let mut handles = Vec::new();
+            for wl in units.chunks(1 + units.len() / n_threads) {
+                handles.push(s.spawn(|| {
+                    let mut n = 0;
+                    wl.iter().for_each(|&(e, new_darts)| {
+                        let mut n_retry = 0;
+                        if map.is_i_free::<2>(e as DartIdType) {
+                            while !process_outer_edge(map, &mut n_retry, e, new_darts)
+                                .is_validated()
+                            {}
+                        } else {
+                            while !process_inner_edge(map, &mut n_retry, e, new_darts)
+                                .is_validated()
+                            {}
+                        }
+                        n += n_retry as u32;
+                    });
+                    n
+                })); // s.spawn
+            } // for wl in workloads
+            handles.into_iter().map(|h| h.join().unwrap()).sum()
+        }) // std::thread::scope
+    }
 }
 
 #[inline]

benches/src/main.rs

@@ -14,6 +14,41 @@ use honeycomb_benches::{
 };
 
 fn main() {
+    #[cfg(feature = "bind-threads")]
+    {
+        use std::sync::Arc;
+
+        use honeycomb_benches::utils::get_proc_list;
+        use hwlocality::{Topology, cpu::binding::CpuBindingFlags};
+        use rayon::ThreadPoolBuilder;
+
+        let builder = ThreadPoolBuilder::new();
+        let topo = Arc::new(Topology::new().unwrap());
+        if let Some(cores) = get_proc_list(&topo) {
+            let mut cores = cores.into_iter().cycle();
+            builder
+                .spawn_handler(|t_builder| {
+                    let topo = topo.clone();
+                    let core = cores.next().expect("E: unreachable"); // due to cycle
+
+                    std::thread::spawn(move || {
+                        // bind
+                        let tid = hwlocality::current_thread_id();
+                        topo.bind_thread_cpu(tid, &core, CpuBindingFlags::empty())
+                            .unwrap();
+                        // work
+                        t_builder.run();
+                    });
+
+                    Ok(())
+                })
+                .build_global()
+                .unwrap();
+        } else {
+            builder.build_global().unwrap()
+        }
+    }
+
     let cli = Cli::parse();
 
     if cli.simple_precision {

benches/src/utils.rs

@@ -3,6 +3,14 @@ use std::fs::File;
 use std::hash::Hasher;
 use std::io::Read;
 
+#[cfg(feature = "bind-threads")]
+use hwlocality::{
+    Topology,
+    cpu::cpuset::CpuSet,
+    object::types::ObjectType,
+    topology::support::{DiscoverySupport, FeatureSupport},
+};
+
 cfg_if::cfg_if! {
     if #[cfg(feature = "_single_precision")] {
         /// Floating-point type alias.
@@ -33,6 +41,144 @@ pub fn hash_file(path: &str) -> Result<u64, std::io::Error> {
     Ok(hasher.finish())
 }
 
+pub const NUM_THREADS_VAR: &str = "RAYON_NUM_THREADS";
+
+pub fn get_num_threads() -> Result<usize, String> {
+    match std::env::var(NUM_THREADS_VAR) {
+        Ok(val) => val.parse::<usize>().map_err(|e| e.to_string()),
+        Err(e) => Err(e.to_string()),
+    }
+}
+
+#[cfg(feature = "bind-threads")]
+#[derive(Debug, Default)]
+pub enum BindingPolicy {
+    /// Disable thread binding.
+    /// Corresponding `RAYON_PROC_BIND_VALUE`: `false`.
+    Disable,
+    /// Enable thread binding & prioritize binding of PUs over cores?.
+    /// Corresponding `RAYON_PROC_BIND_VALUE`: `close`.
+    Close,
+    /// Enable thread binding & prioritize binding across cores over filling PUs?.
+    /// Corresponding `RAYON_PROC_BIND_VALUE`: `spread`. Default value.
+    #[default]
+    Spread,
+}
+
+/// Environment variable controlling the thread-binding policy.
+///
+/// The name of this variable and its possible values reflect the OpenMP equivalents.
+#[cfg(feature = "bind-threads")]
+pub const RAYON_PROC_BIND_VAR: &str = "RAYON_PROC_BIND";
+
+#[cfg(feature = "bind-threads")]
+impl BindingPolicy {
+    fn from_env() -> Self {
+        match std::env::var(RAYON_PROC_BIND_VAR) {
+            Ok(val) => match val.to_lowercase().as_str() {
+                "false" => Self::Disable,
+                "close" => Self::Close,
+                "spread" => Self::Spread,
+                "" => Self::default(),
+                _ => {
+                    eprintln!("W: unrecognized RAYON_PROC_BIND value (!= false | close | spread)");
+                    eprintln!("   continuing with default (spread)");
+                    Self::default()
+                }
+            },
+            Err(e) => {
+                match e {
+                    std::env::VarError::NotPresent => {}
+                    std::env::VarError::NotUnicode(_) => {
+                        eprintln!("W: non-unicode RAYON_PROC_BIND value");
+                        eprintln!("   continuing with default (spread)");
+                    }
+                }
+                Self::default()
+            }
+        }
+    }
+}
+
+#[cfg(feature = "bind-threads")]
+pub fn check_hwloc_support<'a>(topology: &'a Topology) -> Result<(), String> {
+    if !topology.supports(FeatureSupport::discovery, DiscoverySupport::pu_count) {
+        return Err("missing PU reporting support".to_string());
+    }
+    if !topology
+        .feature_support()
+        .cpu_binding()
+        .is_some_and(|s| s.get_thread() && s.set_thread())
+    {
+        return Err("missing binding support".to_string());
+    }
+
+    Ok(())
+}
+
+/// Return a list of bind targets ordered according to the desired policy.
+///
+/// The desired policy is read from an environment variable (see [`RAYON_PROC_BIND_VAR`]). For details on each policy,
+/// see [`BindingPolicy`].
+///
+/// The returned list is used by iterating over a `cycle`d version of it, which corresponds to a round robin logic.
+#[cfg(feature = "bind-threads")]
+pub fn get_proc_list(topology: &Topology) -> Option<Vec<CpuSet>> {
+    let binding_policy = BindingPolicy::from_env();
+    let core_depth = topology
+        .depth_or_below_for_type(ObjectType::Core)
+        .expect("E: unreachable");
+
+    match binding_policy {
+        BindingPolicy::Disable => None,
+        BindingPolicy::Close => {
+            let mut pu_set = Vec::with_capacity(256);
+            topology.objects_at_depth(core_depth).for_each(|c| {
+                let target = c.cpuset().unwrap().clone_target();
+                let w = target.weight();
+                if !(w == Some(1) || w == Some(2)) {
+                    panic!()
+                }
+                target
+                    .iter_set()
+                    .map(CpuSet::from)
+                    .for_each(|t| pu_set.push(t));
+            });
+            Some(pu_set)
+        }
+        BindingPolicy::Spread => {
+            let mut first_pu_set = Vec::with_capacity(128);
+            let mut second_pu_set = Vec::with_capacity(128);
+            topology.objects_at_depth(core_depth).for_each(|c| {
+                let target = c.cpuset().expect("E: unreachable").clone_target();
+                // match required bc some modern CPUs have HT/SMT on only some of their cores :)
+                match target.weight() {
+                    Some(1) => {
+                        // one PU per core -> no HT/SMT
+                        first_pu_set.push(target);
+                    }
+                    Some(2) => {
+                        // two PUs per core -> HT/SMT
+                        let [first_pu, second_pu]: [CpuSet; 2] = target
+                            .iter_set()
+                            .map(CpuSet::from)
+                            .collect::<Vec<_>>()
+                            .try_into()
+                            .expect("E: unreachable");
+                        first_pu_set.push(first_pu);
+                        second_pu_set.push(second_pu);
+                    }
+                    Some(_) | None => {
+                        panic!("E: architecture too cursed")
+                    }
+                }
+            });
+            first_pu_set.append(&mut second_pu_set);
+            Some(first_pu_set)
+        }
+    }
+}
+
 #[cfg(feature = "profiling")]
 pub static mut PERF_FIFO: Option<File> = None;