refactor: make matrix multiply to be concurrency with threads

Author: sword-demon

Cargo.lock

@@ -1,10 +1,11 @@
 [package]
-name = "template"
+name = "concurrency"
 version = "0.1.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 anyhow = "1.0.86"
+oneshot = "0.1.8"
 rand = "0.8.5"

Cargo.toml

@@ -1,6 +1,9 @@
 use anyhow::{Ok, Result};
+use concurrency::Matrix;
 
 fn main() -> Result<()> {
-    println!("f64 default: {}", f64::default());
+    let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
+    let b = Matrix::new([1, 2, 3, 4, 5, 6], 3, 2);
+    println!("a * b = {}", a * b);
     Ok(())
 }

examples/matrix.rs

@@ -1,3 +1,4 @@
 mod matrix;
-
+mod vector;
 pub use matrix::{multiply, Matrix};
+pub use vector::*;

src/lib.rs

@@ -1,9 +1,16 @@
-use anyhow::{anyhow, Result};
 use std::{
     fmt,
     ops::{Add, AddAssign, Mul},
+    sync::mpsc,
+    thread,
 };
 
+use anyhow::{anyhow, Result};
+
+use crate::{dot_product, Vector};
+
+const NUM_THREADS: usize = 4;
+
 // 声明一个矩阵的结构
 
 pub struct Matrix<T> {
@@ -12,30 +19,104 @@ pub struct Matrix<T> {
     col: usize,
 }
 
+pub struct MsgInput<T> {
+    idx: usize,
+    row: Vector<T>,
+    col: Vector<T>,
+}
+
+pub struct MsgOutput<T> {
+    idx: usize,
+    value: T,
+}
+
+pub struct Msg<T> {
+    input: MsgInput<T>,
+    // sender to send the result back
+    sender: oneshot::Sender<MsgOutput<T>>,
+}
+
 pub fn multiply<T>(a: &Matrix<T>, b: &Matrix<T>) -> Result<Matrix<T>>
 where
-    T: Copy + Default + Add<Output = T> + AddAssign + Mul<Output = T>,
+    T: Copy + Default + Add<Output = T> + AddAssign + Mul<Output = T> + Send + 'static,
 {
     if a.col != b.row {
         return Err(anyhow!("Matrix multiply error: a.col != b.row"));
     }
 
+    let senders = (0..NUM_THREADS)
+        .map(|_| {
+            let (tx, rx) = mpsc::channel::<Msg<T>>();
+            // 线程下执行的内容
+            thread::spawn(move || {
+                for msg in rx {
+                    let value = dot_product(msg.input.row, msg.input.col)?;
+                    if let Err(e) = msg.sender.send(MsgOutput {
+                        idx: msg.input.idx,
+                        value,
+                    }) {
+                        eprintln!("send error: {:?}", e);
+                    }
+                }
+                // 这里我犯了一个错误，加了分号导致找了很长时间的 bug
+                Ok::<_, anyhow::Error>(())
+            });
+            tx
+        })
+        .collect::<Vec<_>>();
+
+    // generate 4 threads which receive msg and do dot product
+
     // 让数据有缺省值
-    let mut data = vec![T::default(); a.row * b.col];
+    let matrix_len = a.row * b.col;
+    // 初始化一个空的矩阵
+    let mut data = vec![T::default(); matrix_len];
+    let mut receivers = Vec::with_capacity(matrix_len);
+    // map reduce map phase
     for i in 0..a.row {
         for j in 0..b.col {
-            for k in 0..a.col {
-                data[i * b.col + j] += a.data[i * a.col + k] * b.data[k * b.col + j]
+            let row = Vector::new(&a.data[i * a.col..(i + 1) * a.col]);
+            let col_data = b.data[j..]
+                .iter()
+                .step_by(b.col)
+                .copied()
+                .collect::<Vec<_>>();
+            let col = Vector::new(col_data);
+            let idx = i * b.col + j;
+            let input = MsgInput::new(idx, row, col);
+            let (tx, rx) = oneshot::channel();
+            let msg = Msg::new(input, tx);
+            if let Err(e) = senders[idx % NUM_THREADS].send(msg) {
+                eprintln!("send error: {:?}", e);
             }
+            receivers.push(rx);
         }
     }
+
+    // map reduce: reduce phase
+    for rx in receivers {
+        let output = rx.recv()?;
+        data[output.idx] = output.value;
+    }
     Ok(Matrix {
         data,
         row: a.row,
         col: b.col,
     })
 }
 
+impl<T> MsgInput<T> {
+    pub fn new(idx: usize, row: Vector<T>, col: Vector<T>) -> Self {
+        Self { idx, row, col }
+    }
+}
+
+impl<T> Msg<T> {
+    pub fn new(input: MsgInput<T>, sender: oneshot::Sender<MsgOutput<T>>) -> Self {
+        Self { input, sender }
+    }
+}
+
 impl<T: fmt::Debug> Matrix<T> {
     pub fn new(data: impl Into<Vec<T>>, row: usize, col: usize) -> Self {
         Self {
@@ -79,20 +160,59 @@ where
     }
 }
 
+impl<T> Mul for Matrix<T>
+where
+    T: Copy + Default + Add<Output = T> + AddAssign + Mul<Output = T> + Send + 'static,
+{
+    type Output = Self;
+
+    fn mul(self, rhs: Self) -> Self::Output {
+        multiply(&self, &rhs).expect("Matrix multiply error")
+    }
+}
+
 #[cfg(test)]
 mod tests {
+    use std::vec;
+
     use super::*;
 
     #[test]
     fn test_matrix_multiply() -> Result<()> {
         let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
         let b = Matrix::new([1, 2, 3, 4, 5, 6], 3, 2);
-        let c = multiply(&a, &b)?;
+        let c = a * b;
         assert_eq!(c.col, 2);
         assert_eq!(c.row, 2);
         assert_eq!(c.data, vec![22, 28, 49, 64]);
         assert_eq!(format!("{:?}", c), "Matrix(row=2, col=2, {22 28, 49 64})");
 
         Ok(())
     }
+
+    #[test]
+    fn test_matrix_display() -> Result<()> {
+        let a = Matrix::new([1, 2, 3, 4], 2, 2);
+        let b = Matrix::new([1, 2, 3, 4], 2, 2);
+        let c = a * b;
+        assert_eq!(c.data, vec![7, 10, 15, 22]);
+        assert_eq!(format!("{}", c), "{7 10, 15 22}");
+        Ok(())
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_a_can_not_multiply_b_panic() {
+        let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
+        let b = Matrix::new([1, 2, 3, 4], 2, 2);
+        let _c = a * b;
+    }
+
+    #[test]
+    fn test_a_can_not_multiply_b() {
+        let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
+        let b = Matrix::new([1, 2, 3, 4], 2, 2);
+        let c = multiply(&a, &b);
+        assert!(c.is_err());
+    }
 }

src/matrix.rs

@@ -0,0 +1,51 @@
+use anyhow::{anyhow, Result};
+use std::ops::{Add, AddAssign, Deref, Mul};
+
+pub struct Vector<T> {
+    data: Vec<T>,
+}
+
+impl<T> Deref for Vector<T> {
+    type Target = Vec<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.data
+    }
+}
+
+impl<T> Vector<T> {
+    pub fn new(data: impl Into<Vec<T>>) -> Self {
+        Self { data: data.into() }
+    }
+
+    // pub fn len(&self) -> usize {
+    //     self.data.len()
+    // }
+
+    // pub fn iter(&self) -> std::slice::Iter<T> {
+    //     self.data.iter()
+    // }
+}
+
+// impl<T> Index<usize> for Vector<T> {
+//     type Output = T;
+//     fn index(&self, index: usize) -> &Self::Output {
+//         &self.data[index]
+//     }
+// }
+
+pub fn dot_product<T>(a: Vector<T>, b: Vector<T>) -> Result<T>
+where
+    T: Copy + Default + Add<Output = T> + AddAssign + Mul<Output = T>,
+{
+    if a.len() != b.len() {
+        return Err(anyhow!("Dot product error: a.len != b.len"));
+    }
+
+    let mut sum = T::default();
+    for i in 0..a.len() {
+        sum += a[i] * b[i];
+    }
+
+    Ok(sum)
+}