diff --git a/Cargo.lock b/Cargo.lock index 6a72436..6e0b283 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -325,9 +325,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] @@ -638,9 +638,9 @@ checksum = "6b947ae49db0d222b1dbc6b113ce7248a3fc3a6ca21b696717bfc000ba4484d8" [[package]] name = "memmap2" -version = "0.9.10" +version = "0.9.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" dependencies = [ "libc", ] @@ -752,7 +752,13 @@ dependencies = [ name = "mohu-ops" version = "0.1.0" dependencies = [ + "half", + "mohu-buffer", "mohu-core", + "mohu-dtype", + "mohu-error", + "mohu-testing", + "num-complex", "num-traits", "rayon", "thiserror", diff --git a/crates/mohu-ops/Cargo.toml b/crates/mohu-ops/Cargo.toml index a29ee8a..f5ab3a3 100644 --- a/crates/mohu-ops/Cargo.toml +++ b/crates/mohu-ops/Cargo.toml @@ -12,10 +12,19 @@ categories.workspace = true [dependencies] mohu-core.workspace = true +mohu-buffer.workspace = true +mohu-error.workspace = true +mohu-dtype.workspace = true num-traits.workspace = true rayon.workspace = true thiserror.workspace = true +half.workspace = true +num-complex.workspace = true + +[dev-dependencies] +mohu-testing = { workspace = true } + [package.metadata.cargo-machete] -ignored = ["mohu-core", "num-traits", "rayon", "thiserror"] +ignored = ["mohu-core", "num-traits", "rayon", "thiserror", "half", "num-complex"] \ No newline at end of file diff --git a/crates/mohu-ops/src/lib.rs b/crates/mohu-ops/src/lib.rs index 46515b4..63a0491 100644 --- a/crates/mohu-ops/src/lib.rs +++ b/crates/mohu-ops/src/lib.rs @@ -2,5 +2,6 @@ pub mod arith; pub mod broadcast; pub mod cmp; pub mod logical; +pub mod matmul; pub mod reduce; pub mod unary; diff --git a/crates/mohu-ops/src/matmul.rs b/crates/mohu-ops/src/matmul.rs new file mode 100644 index 0000000..7724af6 --- /dev/null +++ b/crates/mohu-ops/src/matmul.rs @@ -0,0 +1,128 @@ +use mohu_buffer::{buffer::Buffer, layout::Order}; +use mohu_dtype::{ + dispatch_numeric, + dtype::DType, + promote::{CastMode, promote}, + scalar::Scalar, +}; +use mohu_error::{MohuError, MohuResult}; + +pub fn matmul(lhs: &Buffer, rhs: &Buffer) -> MohuResult { + let lhs_shape = lhs.shape().to_vec(); + let rhs_shape = rhs.shape().to_vec(); + + // Normalize every supported case into an effective (m, k, n) plus the + // real output shape the caller should see. + let (m, k, n, out_shape): (usize, usize, usize, Vec) = + match (lhs_shape.len(), rhs_shape.len()) { + // (M, K) x (K, N) -> (M, N) + (2, 2) => { + let (m, k) = (lhs_shape[0], lhs_shape[1]); + let (k2, n) = (rhs_shape[0], rhs_shape[1]); + check_k(k, k2)?; + (m, k, n, vec![m, n]) + }, + // (K,) as (1, K); x (K, N) -> (N,) + (1, 2) => { + let k = lhs_shape[0]; + let (k2, n) = (rhs_shape[0], rhs_shape[1]); + check_k(k, k2)?; + (1, k, n, vec![n]) + }, + // (M, K) x (K,) as (K, 1) -> (M,) + (2, 1) => { + let (m, k) = (lhs_shape[0], lhs_shape[1]); + let k2 = rhs_shape[0]; + check_k(k, k2)?; + (m, k, 1, vec![m]) + }, + // (K,) x (K,) -> scalar, shape [] + (1, 1) => { + let k = lhs_shape[0]; + let k2 = rhs_shape[0]; + check_k(k, k2)?; + (1, k, 1, vec![]) + }, + (a, b) => { + // TODO: confirm the actual MohuError variant name for this — + // couldn't find `pub enum MohuError` in lib.rs via grep, so it's + // likely declared elsewhere (error.rs?) or via a macro. + return Err(MohuError::DimensionMismatch { + expected: 2, + got: a.max(b), + }); + }, + }; + + // Spec: promote dtypes, but integer results must promote to F64. + let promoted = promote(lhs.dtype(), rhs.dtype()); + let out_dtype = if promoted.is_integer() { + DType::F64 + } else { + promoted + }; + + // Cast both operands to out_dtype so the kernel only ever deals with one + // type. This also guarantees a contiguous, row-major layout for the + // index math below, since `cast` allocates fresh Order::C storage + // (and short-circuits to `to_contiguous()` when no dtype change is needed). + // TODO: confirm the right CastMode variant here (Safe/Checked/Lossy — whatever + // the crate calls "just do a numeric cast", as opposed to a bit-reinterpret). + let lhs_cast = lhs.cast(out_dtype, CastMode::Safe)?; + let rhs_cast = rhs.cast(out_dtype, CastMode::Safe)?; + + let mut out = Buffer::alloc(out_dtype, &out_shape, Order::C)?; + + macro_rules! do_matmul { + ($T:ty) => { + matmul_typed::<$T>(&lhs_cast, &rhs_cast, &mut out, m, k, n) + }; + } + dispatch_numeric!(out_dtype, do_matmul)??; + + Ok(out) +} + +#[inline] +fn check_k(k: usize, k2: usize) -> MohuResult<()> { + if k != k2 { + return Err(MohuError::ShapeMismatch { + expected: vec![k], + got: vec![k2], + }); + } + Ok(()) +} + +/// Naive O(M*K*N) matmul kernel over logical (m, k) x (k, n) shapes. +/// Assumes `lhs`, `rhs`, and `out` are contiguous and share dtype `T` +/// (guaranteed by the `cast` calls in `matmul`). +fn matmul_typed( + lhs: &Buffer, + rhs: &Buffer, + out: &mut Buffer, + m: usize, + k: usize, + n: usize, +) -> MohuResult<()> +where + T: Scalar + std::ops::Add + std::ops::Mul, +{ + let lhs_data = lhs.as_slice::()?; + let rhs_data = rhs.as_slice::()?; + let out_data = out.as_mut_slice::()?; + + for i in 0..m { + for j in 0..n { + let mut sum = T::ZERO; + // TODO: replace with BLAS dgemm/sgemm once the linear algebra + // backend is available. + for p in 0..k { + sum = sum + lhs_data[i * k + p] * rhs_data[p * n + j]; + } + out_data[i * n + j] = sum; + } + } + + Ok(()) +} diff --git a/crates/mohu-ops/tests/matmul_tests.rs b/crates/mohu-ops/tests/matmul_tests.rs new file mode 100644 index 0000000..f16d83b --- /dev/null +++ b/crates/mohu-ops/tests/matmul_tests.rs @@ -0,0 +1,107 @@ +use mohu_buffer::buffer::Buffer; +use mohu_ops::matmul::matmul; + +#[test] +fn matmul_2x2() { + let a = Buffer::from_vec(vec![1.0f32, 2.0, 3.0, 4.0]) + .unwrap() + .reshape(&[2, 2]) + .unwrap(); + + let b = Buffer::from_vec(vec![5.0f32, 6.0, 7.0, 8.0]) + .unwrap() + .reshape(&[2, 2]) + .unwrap(); + + let c = matmul(&a, &b).unwrap(); + + assert_eq!(c.shape(), &[2, 2]); + + let out = c.as_slice::().unwrap(); + + mohu_testing::approx::assert_allclose(out, &[19.0, 22.0, 43.0, 50.0]); +} + +#[test] +fn matmul_shape_mismatch() { + let a = Buffer::from_vec(vec![1.0f32; 6]) + .unwrap() + .reshape(&[2, 3]) + .unwrap(); + + let b = Buffer::from_vec(vec![1.0f32; 8]) + .unwrap() + .reshape(&[4, 2]) + .unwrap(); + + assert!(matmul(&a, &b).is_err()); +} + +#[test] +fn matmul_dot_product() { + let a = Buffer::from_vec(vec![1.0f32, 2.0, 3.0]).unwrap(); + let b = Buffer::from_vec(vec![4.0f32, 5.0, 6.0]).unwrap(); + + let c = matmul(&a, &b).unwrap(); + + assert_eq!(c.shape(), &[] as &[usize]); + + let out = c.as_slice::().unwrap(); + + assert_eq!(out, &[32.0]); +} + +#[test] +fn matmul_row_vector_matrix() { + let a = Buffer::from_vec(vec![1.0f32, 2.0]).unwrap(); + + let b = Buffer::from_vec(vec![3.0f32, 4.0, 5.0, 6.0]) + .unwrap() + .reshape(&[2, 2]) + .unwrap(); + + let c = matmul(&a, &b).unwrap(); + + assert_eq!(c.shape(), &[2]); + + let out = c.as_slice::().unwrap(); + + assert_eq!(out, &[13.0, 16.0]); +} + +#[test] +fn matmul_matrix_vector() { + let a = Buffer::from_vec(vec![1.0f32, 2.0, 3.0, 4.0]) + .unwrap() + .reshape(&[2, 2]) + .unwrap(); + + let b = Buffer::from_vec(vec![5.0f32, 6.0]).unwrap(); + + let c = matmul(&a, &b).unwrap(); + + assert_eq!(c.shape(), &[2]); + + let out = c.as_slice::().unwrap(); + + assert_eq!(out, &[17.0, 39.0]); +} + +use mohu_dtype::dtype::DType; + +#[test] +fn matmul_integer_promotes_to_f64() { + let a = Buffer::from_vec(vec![1i32, 2, 3, 4]) + .unwrap() + .reshape(&[2, 2]) + .unwrap(); + + let b = Buffer::from_vec(vec![5i32, 6, 7, 8]) + .unwrap() + .reshape(&[2, 2]) + .unwrap(); + + let c = matmul(&a, &b).unwrap(); + + assert_eq!(c.dtype(), DType::F64); +}