Rust
SIMD matrix
Do parallel matrix multiplication with SIMD.
Final task: SIMD matrix
Write a 4×4 matrix-vector product in a SIMD-friendly form. The classic row-dot approach needs a horizontal reduction for every row. The column approach uses only vertical operations: the result vector is the sum of each column multiplied by the corresponding x component (y += column_j * x[j], "axpy"):
// y = A·x, A in column order: a[j] = column j
fn matvec(a: [[f32; 4]; 4], x: [f32; 4]) -> [f32; 4] {
let mut y = [0.0f32; 4];
for j in 0..4 {
for i in 0..4 { // a 4-lane vertical multiply-add
y[i] += a[j][i] * x[j];
}
}
y
}
The inner loop maps directly onto a single 4-lane multiply-add instruction; no horizontal operation is needed.
Verification
An optimized kernel must always be compared with a simple reference implementation. A reference giving the same result with a row-ordered matrix:
fn matvec_ref(satirlar: [[f32; 4]; 4], x: [f32; 4]) -> [f32; 4] {
satirlar.map(|r| r.iter().zip(&x).map(|(a, b)| a * b).sum())
}
// converting to column order: the transpose
fn devrik(m: [[f32; 4]; 4]) -> [[f32; 4]; 4] {
std::array::from_fn(|j| std::array::from_fn(|i| m[i][j]))
}
Exercise Matrix-vector
Write the matvec, matvec_ref and devrik functions above. For the row-ordered [[1,2,0,0],[0,1,0,0],[0,0,2,0],[0,0,0,3]] and x = [1, 1, 1, 1], print the result of matvec(devrik(A), x) with {:?} and its equality with the reference, separated by a space.
Show example solution
fn matvec(a: [[f32; 4]; 4], x: [f32; 4]) -> [f32; 4] {
let mut y = [0.0f32; 4];
for j in 0..4 {
for i in 0..4 {
y[i] += a[j][i] * x[j];
}
}
y
}
fn matvec_ref(satirlar: [[f32; 4]; 4], x: [f32; 4]) -> [f32; 4] {
satirlar.map(|r| r.iter().zip(&x).map(|(a, b)| a * b).sum())
}
fn devrik(m: [[f32; 4]; 4]) -> [[f32; 4]; 4] {
std::array::from_fn(|j| std::array::from_fn(|i| m[i][j]))
}
fn main() {
let a = [
[1.0, 2.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 2.0, 0.0],
[0.0, 0.0, 0.0, 3.0],
];
let x = [1.0; 4];
let y = matvec(devrik(a), x);
println!("{:?} {}", y, y == matvec_ref(a, x));
}
Test yourself
-
What is the advantage of column-order storage for this product?
- Less memory
- It's easier to write
- No advantage
- Each column is contiguous in memory; it fits into a vector with a single load
Show answer
Each column is contiguous in memory; it fits into a vector with a single load — Choosing the data layout for the algorithm determines most of the SIMD performance.
-
Vertical multiply-add
for j in 0..4 {
for i in 0..4 {
y[i] += a[j][i] * x[____];
}
}
Show answer
for j in 0..4 {
for i in 0..4 {
y[i] += a[j][i] * x[j];
}
}
-
Approach and property
- Row-dot
- Needs horizontal reduction
- Column axpy
- Only vertical
- Referans uygulama
- Verification
- Devrik
- Layout conversion
Next lesson: no_std runtime
