Implement vectorized matmul

1 files changed, 147 insertions, 5 deletions

diff --git a/src/matrix.cpp b/src/matrix.cpp
index c6888a4..1faab25 100644
--- a/src/matrix.cpp
+++ b/src/matrix.cpp
@@ -2,9 +2,19 @@
 #include "linalg_error.hpp"
 
 #include <algorithm>
+#include <cstddef>
 #include <sstream>
 #include <stdexcept>
 
+#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && \
+    (defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__))
+#include <immintrin.h>
+#endif
+
+#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__ARM_NEON) && defined(__aarch64__) && !defined(__clangd__)
+#include <arm_neon.h>
+#endif
+
 namespace linalg {
 
 namespace {
@@ -17,6 +27,132 @@ void check_same_shape(const Matrix& lhs, const Matrix& rhs, const char* operatio
     }
 }
 
+double dot_product_scalar(const double* lhs, const double* rhs, std::size_t count) {
+    double sum = 0.0;
+    for (std::size_t i = 0; i < count; ++i) {
+        sum += lhs[i] * rhs[i];
+    }
+    return sum;
+}
+
+#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__AVX512F__)
+double horizontal_sum(__m512d values) {
+    alignas(64) double lanes[8];
+    _mm512_store_pd(lanes, values);
+    double sum = 0.0;
+    for (double lane : lanes) {
+        sum += lane;
+    }
+    return sum;
+}
+
+double dot_product_avx512(const double* lhs, const double* rhs, std::size_t count) {
+    std::size_t i = 0;
+    __m512d acc0 = _mm512_setzero_pd();
+    __m512d acc1 = _mm512_setzero_pd();
+
+    for (; i + 15 < count; i += 16) {
+        const __m512d lhs0 = _mm512_loadu_pd(lhs + i);
+        const __m512d rhs0 = _mm512_loadu_pd(rhs + i);
+        const __m512d lhs1 = _mm512_loadu_pd(lhs + i + 8);
+        const __m512d rhs1 = _mm512_loadu_pd(rhs + i + 8);
+
+        acc0 = _mm512_add_pd(acc0, _mm512_mul_pd(lhs0, rhs0));
+        acc1 = _mm512_add_pd(acc1, _mm512_mul_pd(lhs1, rhs1));
+    }
+
+    return horizontal_sum(acc0) + horizontal_sum(acc1) +
+           dot_product_scalar(lhs + i, rhs + i, count - i);
+}
+#endif
+
+#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__AVX2__)
+double horizontal_sum(__m256d values) {
+    alignas(32) double lanes[4];
+    _mm256_store_pd(lanes, values);
+    return lanes[0] + lanes[1] + lanes[2] + lanes[3];
+}
+
+double dot_product_avx2(const double* lhs, const double* rhs, std::size_t count) {
+    std::size_t i = 0;
+    __m256d acc0 = _mm256_setzero_pd();
+    __m256d acc1 = _mm256_setzero_pd();
+
+    for (; i + 7 < count; i += 8) {
+        const __m256d lhs0 = _mm256_loadu_pd(lhs + i);
+        const __m256d rhs0 = _mm256_loadu_pd(rhs + i);
+        const __m256d lhs1 = _mm256_loadu_pd(lhs + i + 4);
+        const __m256d rhs1 = _mm256_loadu_pd(rhs + i + 4);
+
+        acc0 = _mm256_add_pd(acc0, _mm256_mul_pd(lhs0, rhs0));
+        acc1 = _mm256_add_pd(acc1, _mm256_mul_pd(lhs1, rhs1));
+    }
+
+    return horizontal_sum(acc0) + horizontal_sum(acc1) +
+           dot_product_scalar(lhs + i, rhs + i, count - i);
+}
+#endif
+
+#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__AVX__) && !defined(__AVX2__)
+double horizontal_sum(__m256d values) {
+    alignas(32) double lanes[4];
+    _mm256_store_pd(lanes, values);
+    return lanes[0] + lanes[1] + lanes[2] + lanes[3];
+}
+
+double dot_product_avx(const double* lhs, const double* rhs, std::size_t count) {
+    std::size_t i = 0;
+    __m256d acc = _mm256_setzero_pd();
+
+    for (; i + 3 < count; i += 4) {
+        const __m256d lhs_values = _mm256_loadu_pd(lhs + i);
+        const __m256d rhs_values = _mm256_loadu_pd(rhs + i);
+        acc = _mm256_add_pd(acc, _mm256_mul_pd(lhs_values, rhs_values));
+    }
+
+    return horizontal_sum(acc) + dot_product_scalar(lhs + i, rhs + i, count - i);
+}
+#endif
+
+#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__ARM_NEON) && defined(__aarch64__) && !defined(__clangd__)
+double horizontal_sum(float64x2_t values) {
+    return vgetq_lane_f64(values, 0) + vgetq_lane_f64(values, 1);
+}
+
+double dot_product_neon(const double* lhs, const double* rhs, std::size_t count) {
+    std::size_t i = 0;
+    float64x2_t acc0 = vdupq_n_f64(0.0);
+    float64x2_t acc1 = vdupq_n_f64(0.0);
+
+    for (; i + 3 < count; i += 4) {
+        const float64x2_t lhs0 = vld1q_f64(lhs + i);
+        const float64x2_t rhs0 = vld1q_f64(rhs + i);
+        const float64x2_t lhs1 = vld1q_f64(lhs + i + 2);
+        const float64x2_t rhs1 = vld1q_f64(rhs + i + 2);
+
+        acc0 = vaddq_f64(acc0, vmulq_f64(lhs0, rhs0));
+        acc1 = vaddq_f64(acc1, vmulq_f64(lhs1, rhs1));
+    }
+
+    return horizontal_sum(acc0) + horizontal_sum(acc1) +
+           dot_product_scalar(lhs + i, rhs + i, count - i);
+}
+#endif
+
+double dot_product_simd(const double* lhs, const double* rhs, std::size_t count) {
+#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__AVX512F__)
+    return dot_product_avx512(lhs, rhs, count);
+#elif !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__AVX2__)
+    return dot_product_avx2(lhs, rhs, count);
+#elif !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__AVX__) && !defined(__AVX2__)
+    return dot_product_avx(lhs, rhs, count);
+#elif !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__ARM_NEON) && defined(__aarch64__) && !defined(__clangd__)
+    return dot_product_neon(lhs, rhs, count);
+#else
+    return dot_product_scalar(lhs, rhs, count);
+#endif
+}
+
 }  // namespace
 
 Matrix::Matrix(std::size_t rows, std::size_t cols)
@@ -144,13 +280,19 @@ Matrix operator*(const Matrix& lhs, const Matrix& rhs) {
         throw DimensionMismatchError(oss.str());
     }
 
+    const Matrix rhs_transposed = transpose(rhs);
     Matrix result(lhs.rows(), rhs.cols());
+
+    const std::size_t inner_dim = lhs.cols();
+    const double* lhs_data = lhs.data();
+    const double* rhs_t_data = rhs_transposed.data();
+    double* result_data = result.data();
+
     for (std::size_t i = 0; i < lhs.rows(); ++i) {
-        for (std::size_t k = 0; k < lhs.cols(); ++k) {
-            const double lhs_ik = lhs(i, k);
-            for (std::size_t j = 0; j < rhs.cols(); ++j) {
-                result(i, j) += lhs_ik * rhs(k, j);
-            }
+        const double* lhs_row = lhs_data + i * inner_dim;
+        for (std::size_t j = 0; j < rhs.cols(); ++j) {
+            const double* rhs_column = rhs_t_data + j * inner_dim;
+            result_data[i * rhs.cols() + j] = dot_product_simd(lhs_row, rhs_column, inner_dim);
         }
     }
     return result;