zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Implemented Karatsuba's Algorithm for complex matmul and pre-commit them

Browse Source
This commit is contained in:
thesuryash
2025-05-23 08:30:47 -04:00
parent 1d9076c24b
commit f859e75f4f
1 changed files with 16 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
mlx/ops.cpp
View File

@@ -2863,30 +2863,30 @@ array matmul(
throw std::invalid_argument(msg.str()); throw std::invalid_argument(msg.str());
} }
// complex matmul using Karatsuba's Algorithm // complex matmul using Karatsuba's Algorithm
if (a.dtype() == complex64 && b.dtype() == complex64) { if (a.dtype() == complex64 && b.dtype() == complex64) {
// Extract real and imaginary parts // Extract real and imaginary parts
auto a_real = real(a, s); auto a_real = real(a, s);
auto a_imag = imag(a, s); auto a_imag = imag(a, s);
auto b_real = real(b, s); auto b_real = real(b, s);
auto b_imag = imag(b, s); auto b_imag = imag(b, s);
// Compute real and imaginary components of the result // Compute real and imaginary components of the result
auto m1 = matmul(a_real, b_real, s); auto m1 = matmul(a_real, b_real, s);
auto m2 = matmul(a_imag, b_imag, s); auto m2 = matmul(a_imag, b_imag, s);
auto m3 = matmul(add(a_real, a_imag, s), add(b_real, b_imag, s), s); auto m3 = matmul(add(a_real, a_imag, s), add(b_real, b_imag, s), s);
auto c_real = subtract(m1, m2, s); auto c_real = subtract(m1, m2, s);
auto c_imag = subtract(m3, add(m1, m2, s), s); auto c_imag = subtract(m3, add(m1, m2, s), s);
return add(c_real, multiply(array(complex64_t{0, 1}, complex64), c_imag, s), s); return add(
} c_real, multiply(array(complex64_t{0, 1}, complex64), c_imag, s), s);
}
// Type promotion // Type promotion
auto out_type = promote_types(a.dtype(), b.dtype()); auto out_type = promote_types(a.dtype(), b.dtype());
if (!issubdtype(out_type, floating)) { if (!issubdtype(out_type, floating)) {
std::ostringstream msg; std::ostringstream msg;
msg << "[matmul] Only real floating point types are supported but " msg << "[matmul] Only real floating point types are supported but "