Scaling with rayon

Hi everyone!

I'm using rust for doing some computationally intense statistics with images (astronomy). I wrote a rust program that uses rayon for parallelization. I've ran my program on two boxes with the following specs:

1. 4x Intel Xeon E5-4620v2 (32 cores total) with 505GB ECC DDR4 ram running at base speed
2. 1x Intel i5-1135G7 (4 cores) with 32GB DDR4 ram @3200MHz

Reducing a single image takes \2min on my laptop and \43min (!!!) on the server. I have no clue why my code is so much slower on the more powerful computer. For context, the images are entirely loaded into RAM before the timed computation starts, so there is no network/IO overhead for the server in this comparison.

I guess I should talk a little bit about my code. The slow bit is a watershed-like transform of each image. It consists of an outer, non-rayon for loop that increases the water-level from 0 to 255:

for water_level in (0..u8::MAX).step_by(FILL_STEP as usize) {
    //(...)
}

Each time I increase the water level I filter out all flooded pixels:

let flooded: Vec<(usize, usize)> = img
    .indexed_iter()
    .par_bridge()
    //Filter flooded pixels
    .filter(|&(_idx, img_val)| *img_val <= water_level)
    //Filter uncoloured pixel
    .filter(|&(idx, _img_val)| *out.get(idx).unwrap() == UNCOLOURED)
    .map(|(idx, _val)| idx)
    .collect();

Next, I loop trying to find flooded pixels that are adjacent to already coloured pixels until there are no pixels to color left:

loop {
    let fill_this_tick: Vec<((usize, usize), usize)> = (omitted);

    if fill_this_tick.len() == 0 {
        //We have to raise the water level
        break;
    } else {
        //We colour the pixels <3
        for (idx, colour) in fill_this_tick {
            out[idx] = colour;
        }
    }
}

The last component is how I determine which pixels to colour. I do this by iterating with rayon over all the uncoloured flooded pixels that I found in the first step, and checking if they touch a coloured pixel or not.

let fill_this_tick: Vec<((usize, usize), usize)> = flooded
   .par_iter()
   //if a neighbouring pixel is coloured, we take its colour
   //filter_map takes care of out-of-bounds pixels
   .filter_map(|idx1| {
       let coloured_neighbours = neighbours_4con(idx1)
       .into_iter()
       //Ignore out-of-bounds
       .filter_map(|idx2| out.get(idx2))
       //Ignore uncoloured
       .filter(|col| **col != UNCOLOURED)
       //Get rid of filter pointer
       .map(|x| *x)
       .collect::<Vec<usize>>();
    //we distinguish three cases:
    // (1) all neighbours have same colour (island)
    // (2) some neighbours have different colours (edge)
    // (3) there are no coloured neighbours (do nothing)
    if coloured_neighbours.len() == 0 || *out.get(*idx1).unwrap() != UNCOLOURED {
        //Case (3)
        return None;
    }
    let colour = unsafe { *coloured_neighbours.get_unchecked(0) };
    if coloured_neighbours.iter().all(|&x| x == colour) {
        //Case (1)
        Some((*idx1, colour))
    } else {
        //Case (2), pick a random item from the list
        let rand_idx = rand::thread_rng().gen_range(0..coloured_neighbours.len());
        Some((*idx1, unsafe { *coloured_neighbours.get_unchecked(rand_idx) }))
     }
    })
    .collect();

If you have any ideas/directions for how I could improve the scalability of my code, I'd love to hear them!

Thanks in advance!