Spawning max_num_threads, finding the first

[expenses]

Hi, I have a parallel search problem (Wave Function Collapse) that I'm using parallel iterators for:

pub fn collapse_all_reset_on_contradiction_par(&mut self, rng: &mut SmallRng) -> u32 {
        let stop_flag = Arc::new(AtomicBool::new(false));
        let (wfc, iterations) = (0..rayon::max_num_threads())
            .map(|_| rng.gen::<u64>())
            .par_bridge()
            .map(|seed| (self.clone(), SmallRng::seed_from_u64(seed)))
            .filter_map(|(mut wfc, mut rng)| {
                let mut attempts = 1;
                while let Some((index, tile)) = wfc.find_lowest_entropy(&mut rng) {
                    if stop_flag.load(Ordering::Relaxed) {
                        return None;
                    }
                    if wfc.collapse(index, tile) {
                        if attempts % 1000 == 0 {
                            println!("{}", attempts);
                        }
                        wfc.reset();
                        attempts += 1
                    }
                }

                stop_flag.store(true, Ordering::Relaxed);
                Some((wfc, attempts))
            })
            .find_any(|_| true)
            .unwrap();

        *self = wfc;
        iterations
    }

I really don't like that I need the stop_flag AtomicBool in order to control each thread. Commenting out stop_flag.store(true, Ordering::Relaxed); leads to a situation where the iterator only finishes after every thread has finished; the opposite of what I want! Is there a more elegant way to do this?

[cuviper]

• max_num_threads() is a static property of the internal thread bookkeeping, but you wouldn't actually want that many threads if your system doesn't have that many CPUs. But anyway, the way you're calling this doesn't have any effect on the actual number of threads that will be created in rayon's pool -- it just uses std::thread::available_parallelism() if not otherwise specified in a ThreadPoolBuilder. I think you should just use current_num_threads() instead.

• par_bridge often feels like a waste -- maybe you can just generate the seeded rngs into a quick Vec and then into_par_iter on that? That adds an allocation, but it will avoid the runtime overhead in the bridge.

• Instead of filter_map, can you generate steps with .flat_map_iter? Something like:

fn collapse_all_reset_on_contradiction_par(&mut self, rng: &mut SmallRng) -> u32 {
    let rngs: Vec<_> = (0..rayon::current_num_threads())
        .map(|_| SmallRng::seed_from_u64(rng.gen()))
        .collect();
    let (wfc, iterations) = rngs
        .into_par_iter()
        .flat_map_iter(|mut rng| {
            let mut opt_wfc = Some(self.clone());
            let mut attempts = 1;
            std::iter::from_fn(move || {
                // return None after we already consumed our wfc
                let mut wfc = opt_wfc.take()?;
                if let Some((index, tile)) = wfc.find_lowest_entropy(&mut rng) {
                    if wfc.collapse(index, tile) {
                        if attempts % 1000 == 0 {
                            println!("{}", attempts);
                        }
                        wfc.reset();
                        attempts += 1;
                    }
                    // put it back for the next `find_lowest_entropy` call
                    opt_wfc = Some(wfc);
                    // returning a `None` item gives `find_map_any` a chance to end
                    // this search if a different thread found something
                    Some(None)
                } else {
                    // we're done here, return our wfc!
                    Some(Some((wfc, attempts)))
                }
            })
        })
        .find_map_any(std::convert::identity)
        .unwrap();
    *self = wfc;
    iterations
}

(It's possible that end condition is wrong, but I tried to replicate what I saw in your original code...)

[expenses]

Hmm.. that solution seems a bit too complicated to me, but thanks for the other suggestions. I'm pretty happy with this function:

fn find_any_with_early_stop<
    T,
    O: Send,
    I: IntoParallelIterator<Item = T>,
    F: Sync + Fn(T, &AtomicBool) -> Option<O>,
>(
    iterator: I,
    func: F,
) -> Option<O> {
    let stop_flag = AtomicBool::new(false);
    iterator.into_par_iter().find_map_any(|item| {
        func(item, &stop_flag).inspect(|_| stop_flag.store(true, Ordering::Relaxed))
    })
}