Parallelisation of sparse matrix creation from Vec<CsVec<usize>>

[newtonharry]

Hi,

I'm currently trying to parallelise this code using rayon, but I'm unsure of the approach to take here:

fn convert_to_sparse_matrix(
    matrix: &Vec<CsVec<usize>>,
) -> (Vec<usize>, Vec<usize>, Vec<usize>, (usize, usize)) {
    let num_cols = matrix.len();
    let rows = matrix[0].dim();
    let mut data = Vec::new();
    let mut indices = Vec::new();
    let mut indptr = vec![0];

    for col in matrix.iter() {
        let col_indices = col.indices();
        let col_data = col.data();
        data.extend_from_slice(col_data);
        indices.extend_from_slice(col_indices);
        indptr.push(indptr.last().unwrap() + col.nnz());
    }

    return (data, indices, indptr, (rows, num_cols));
}

fn par_convert_to_sparse_matrix(
    matrix: &Vec<CsVec<usize>>,
) -> (Vec<usize>, Vec<usize>, Vec<usize>, (usize, usize)) {
    let num_cols = matrix.len();
    let rows = matrix[0].dim();
    let chunk_size = 500000; // Size of each chunk for parallel processing
    let num_chunks = (num_cols + chunk_size - 1) / chunk_size;

    // Pre-allocate the final vectors with the necessary capacity
    let total_nnz: usize = matrix.iter().map(|col| col.nnz()).sum();
    let data = Arc::new(Mutex::new(vec![0; total_nnz]));
    let indices = Arc::new(Mutex::new(vec![0; total_nnz]));
    let indptr = Arc::new(Mutex::new(vec![0; num_cols + 1]));

    // Process chunks in parallel
    (0..num_chunks).into_par_iter().for_each(|chunk| {
        let start_col = chunk * chunk_size;
        let end_col = (start_col + chunk_size).min(num_cols);
        let mut local_indptr = vec![0; end_col - start_col + 1];
        let mut local_data = Vec::with_capacity(total_nnz / num_chunks);
        let mut local_indices = Vec::with_capacity(total_nnz / num_chunks);

        for (i, col) in matrix[start_col..end_col].iter().enumerate() {
            let col_indices = col.indices();
            let col_data = col.data();
            local_data.extend_from_slice(col_data);
            local_indices.extend_from_slice(col_indices);
            local_indptr[i + 1] = local_indptr[i] + col.nnz();
        }

        // Lock the global vectors to safely copy the data
        let mut data = data.lock().unwrap();
        let mut indices = indices.lock().unwrap();
        let mut indptr = indptr.lock().unwrap();

        let global_start = indptr[start_col];
        let global_end = global_start + local_data.len();

        data[global_start..global_end].copy_from_slice(&local_data);
        indices[global_start..global_end].copy_from_slice(&local_indices);

        for (i, &local_ptr) in local_indptr.iter().enumerate() {
            indptr[start_col + i + 1] = global_start + local_ptr;
        }
    });

    let data = Arc::try_unwrap(data)
        .expect("Failed to unwrap Arc")
        .into_inner()
        .unwrap();
    let indices = Arc::try_unwrap(indices)
        .expect("Failed to unwrap Arc")
        .into_inner()
        .unwrap();
    let indptr = Arc::try_unwrap(indptr)
        .expect("Failed to unwrap Arc")
        .into_inner()
        .unwrap();

    (data, indices, indptr, (rows, num_cols))
}

Currently the parallelised function compiles but any thoughts on how to do this more effectively?

Thanks!

[mulimoen]

I would think you'd be memory bound and that parallelisation is not going to help much.