Code shared from the Rust Playground

playground.rs

use itertools::*;
use rayon::prelude::*;

type Witness = Vec<Vec<i32>>;

fn main() {
    let challenges = [vec![1, 2, 3], vec![4, 5, 6]];

    // Indexes of each matrix is [column][row]
    let witnesses: Vec<Witness> = vec![
        vec![vec![3, 5, 5], vec![2, 6, 8], vec![9, 9, 6]],
        vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]],
        vec![vec![4, 5, 6], vec![1, 2, 3], vec![7, 8, 9]],
    ];

    let columns_count = witnesses[0].len();
    let rows_count = witnesses[0][0].len();

    // TODO Re-use accumulator
    let empty_witness = vec![vec![0; rows_count]; columns_count];

    let mut result_matrix_by_challenges = vec![empty_witness; challenges.len()];

    iproduct!(0..columns_count, 0..rows_count)
        .map(|(col, row)| {
            witnesses
                .iter()
                .zip(challenges.iter().map(|m| m.iter().copied()).multi_product())
                .fold(
                    vec![0; challenges.len()].into_boxed_slice(),
                    // every element of this collection - one cell for each multiplier
                    |mut result, (matrix, multiplier)| {
                        result
                            .iter_mut()
                            .zip(multiplier.iter())
                            .for_each(|(res, cell)| {
                                *res += cell * matrix[col][row];
                            });

                        result
                    },
                )
        })
        .zip(
            // Here we take an iterator that on each iteration returns [column][row] elements for
            // each matrix for its multiplier
            //
            // next -> vec![matrices[0][col][row], matrices[1][col][row], ... matrices[multiplier_len][col][row]]
            result_matrix_by_challenges
                .iter_mut()
                .map(|matrix| matrix.iter_mut().flat_map(|col| col.iter_mut()))
                .multi_product(),
        )
        .par_bridge()
        .for_each(|(elements, mut results)| {
            results
                .iter_mut()
                .zip(elements.iter())
                .for_each(|(result, cell)| **result = *cell);
        });
}

struct MultiProduct<I: Iterator> {
    iters: Box<[I]>,
}
impl<I: Iterator> Iterator for MultiProduct<I> {
    type Item = Box<[I::Item]>;

    fn next(&mut self) -> Option<Self::Item> {
        self.iters.iter_mut().map(|i| i.next()).collect()
    }
}
trait MultiCartesianProduct: Iterator + Sized
where
    <Self as Iterator>::Item: Iterator + Sized,
{
    fn multi_product(self) -> MultiProduct<Self::Item> {
        MultiProduct {
            iters: self.collect(),
        }
    }
}

impl<I: Iterator + Sized> MultiCartesianProduct for I where
    <Self as Iterator>::Item: Iterator + Sized
{
}