Function blocks 

fn blocks(
    rows: usize,
    cols: usize,
    depth: usize,
) -> impl Iterator<Item = (usize, usize, usize)>

Examples found in repository

crates/competitive/src/num/mint/simd_matrix.rs (line 250)

    pub unsafe fn matrix_product_avx2(
        a: &[Vec<Self>],
        b: &[Vec<Self>],
        scale: u32,
    ) -> Vec<Vec<Self>> {
        let (n, m, p) = (a.len(), b.len(), b.first().map_or(0, Vec::len));
        assert!(a.iter().all(|row| row.len() == m));
        assert!(b.iter().all(|row| row.len() == p));
        let modulus = M::get_mod();
        let alignment = if n.min(m).min(p) <= 64 { 8 } else { 32 };
        let (nn, mm, pp) = (
            n.div_ceil(alignment) * alignment,
            m.div_ceil(alignment) * alignment,
            p.div_ceil(alignment) * alignment,
        );
        let mut depth = 0;
        let (mut x, mut y, mut z) = (nn, mm, pp);
        while x.min(y).min(z) > 64 && x % 16 == 0 && y % 16 == 0 && z % 16 == 0 {
            depth += 1;
            x /= 2;
            y /= 2;
            z /= 2;
        }
        let entries = nn * mm + mm * pp + nn * pp;
        // A recursive level uses one quarter of its parent's storage; siblings reuse it.
        let mut data = if entries + entries / 3 >= 1 << 20 {
            let mut data = Vec::with_capacity(entries + entries / 3);
            advise_huge_pages(&mut data);
            data.resize(entries + entries / 3, 0u32);
            data
        } else {
            vec![0u32; entries + entries / 3]
        };
        let quotient = (((scale as u64) << 32) / modulus as u64) as u32;
        let mut inverse = 1u32;
        for _ in 0..5 {
            inverse = inverse.wrapping_mul(2u32.wrapping_sub(modulus.wrapping_mul(inverse)));
        }
        let inverse = inverse.wrapping_neg();
        for (offset, row, col) in blocks(nn, mm, depth) {
            let (nr, nc) = (nn >> depth, mm >> depth);
            for i in row..(row + nr).min(n) {
                // SAFETY: MInt is transparent over u32. Reading raw words preserves Montgomery encoding.
                let values: &[u32] = unsafe { std::slice::from_raw_parts(a[i].as_ptr().cast(), m) };
                for j in col..(col + nc).min(m) {
                    let x = values[j];
                    let q = ((x as u64 * quotient as u64) >> 32) as u32;
                    let x = x.wrapping_mul(scale).wrapping_sub(q.wrapping_mul(modulus));
                    data[offset + (i - row) * nc + j - col] = x.min(x.wrapping_sub(modulus));
                }
            }
        }
        for (offset, row, col) in blocks(mm, pp, depth) {
            let (nr, nc) = (mm >> depth, pp >> depth);
            for i in row..(row + nr).min(m) {
                // SAFETY: MInt is transparent over u32 and row lengths were checked above.
                let values: &[u32] = unsafe { std::slice::from_raw_parts(b[i].as_ptr().cast(), p) };
                for j in col..(col + nc).min(p) {
                    data[nn * mm + offset + (i - row) * nc + j - col] = values[j];
                }
            }
        }
        let kernel = Kernel {
            modulus,
            inverse,
            avx512: avx512_enabled() && is_x86_feature_detected!("avx512f"),
        };
        // SAFETY: padding keeps every leaf dimension divisible by eight. The three matrices
        // and the geometric scratch space are disjoint parts of the allocated buffer.
        unsafe {
            let ptr = data.as_mut_ptr();
            multiply(
                ptr,
                ptr.add(nn * mm),
                ptr.add(nn * mm + mm * pp),
                (nn, mm, pp),
                ptr.add(entries),
                &kernel,
            );
        }
        let mut result = vec![vec![MInt::new_unchecked(M::mod_zero()); p]; n];
        for (offset, row, col) in blocks(nn, pp, depth) {
            let (nr, nc) = (nn >> depth, pp >> depth);
            for i in row..(row + nr).min(n) {
                for j in col..(col + nc).min(p) {
                    result[i][j] = MInt::new_unchecked(
                        data[nn * mm + mm * pp + offset + (i - row) * nc + j - col],
                    );
                }
            }
        }
        result
    }