Implement vectorized matmul

4 files changed, 202 insertions, 5 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1c854de..c0a79d5 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -9,6 +9,8 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 option(LINEAR_ALGEBRA_BUILD_TESTS "Build unit tests" ON)
 option(LINEAR_ALGEBRA_BUILD_EXAMPLES "Build example programs" ON)
+set(LINEAR_ALGEBRA_SIMD "AUTO" CACHE STRING "SIMD backend for matmul: AUTO, NONE, AVX, AVX2, AVX512")
+set_property(CACHE LINEAR_ALGEBRA_SIMD PROPERTY STRINGS AUTO NONE AVX AVX2 AVX512)
 
 add_library(linear_algebra
     src/vector.cpp
@@ -27,8 +29,26 @@ target_compile_features(linear_algebra PUBLIC cxx_std_20)
 
 if(CMAKE_CXX_COMPILER_ID MATCHES "Clang|GNU")
     target_compile_options(linear_algebra PRIVATE -Wall -Wextra -Wpedantic -Wconversion)
+
+    if(LINEAR_ALGEBRA_SIMD STREQUAL "AVX")
+        target_compile_options(linear_algebra PRIVATE -mavx)
+    elseif(LINEAR_ALGEBRA_SIMD STREQUAL "AVX2")
+        target_compile_options(linear_algebra PRIVATE -mavx2)
+    elseif(LINEAR_ALGEBRA_SIMD STREQUAL "AVX512")
+        target_compile_options(linear_algebra PRIVATE -mavx512f)
+    elseif(LINEAR_ALGEBRA_SIMD STREQUAL "NONE")
+        target_compile_definitions(linear_algebra PRIVATE LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL=1)
+    elseif(LINEAR_ALGEBRA_SIMD STREQUAL "AUTO")
+    else()
+        message(FATAL_ERROR "Unsupported LINEAR_ALGEBRA_SIMD value: ${LINEAR_ALGEBRA_SIMD}")
+    endif()
 elseif(MSVC)
     target_compile_options(linear_algebra PRIVATE /W4 /permissive-)
+    if(LINEAR_ALGEBRA_SIMD STREQUAL "NONE")
+        target_compile_definitions(linear_algebra PRIVATE LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL=1)
+    elseif(NOT LINEAR_ALGEBRA_SIMD STREQUAL "AUTO")
+        message(WARNING "LINEAR_ALGEBRA_SIMD explicit x86 flags are only wired for Clang/GNU right now.")
+    endif()
 endif()
 
 if(LINEAR_ALGEBRA_BUILD_EXAMPLES)
diff --git a/README.md b/README.md
index 6f49640..1cd42ec 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,8 @@
 
 This repo contains a small C++ dense numerical linear algebra library for `double`, with a companion experiments directory for evaluating performance.
 
+The current matmul uses a vectorized dot-product kernel. The implementation supports compile-time SIMD backends for AVX, AVX2, AVX512, NEON (AArch64/ARM64).
+
 ## Build
 
 ```bash
@@ -9,6 +11,14 @@ cmake -S . -B build
 cmake --build build
 ```
 
+On x86, you can explicitly choose a matmul SIMD target at configure time:
+
+```bash
+cmake -S . -B build -DLINEAR_ALGEBRA_SIMD=AVX2
+```
+
+Valid values are `AUTO`, `NONE`, `AVX`, `AVX2`, and `AVX512`. `AUTO` uses the compiler's current target. `NONE` forces the scalar fallback.
+
 ## Run tests
 
 ```bash
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;
diff --git a/tests/test_matrix.cpp b/tests/test_matrix.cpp
index 4cfc783..fcd3d27 100644
--- a/tests/test_matrix.cpp
+++ b/tests/test_matrix.cpp
@@ -174,3 +174,28 @@ TEST_CASE("Matrix-matrix multiply handles identity and shape checks", "[matrix]"
     const Matrix incompatible(4, 1);
     CHECK_THROWS_AS(lhs * incompatible, DimensionMismatchError);
 }
+
+TEST_CASE("Matrix-matrix multiply handles SIMD tail dimensions", "[matrix]") {
+    const Matrix lhs{
+        {1.0, 2.0, 3.0, 4.0, 5.0},
+        {6.0, 7.0, 8.0, 9.0, 10.0}
+    };
+    const Matrix rhs{
+        {1.0, 0.0, 2.0},
+        {0.0, 1.0, 3.0},
+        {1.0, 1.0, 4.0},
+        {0.0, 2.0, 5.0},
+        {1.0, 0.0, 6.0}
+    };
+
+    const Matrix product = lhs * rhs;
+    REQUIRE(product.rows() == 2);
+    REQUIRE(product.cols() == 3);
+
+    CHECK(product(0, 0) == Catch::Approx(9.0));
+    CHECK(product(0, 1) == Catch::Approx(13.0));
+    CHECK(product(0, 2) == Catch::Approx(70.0));
+    CHECK(product(1, 0) == Catch::Approx(24.0));
+    CHECK(product(1, 1) == Catch::Approx(33.0));
+    CHECK(product(1, 2) == Catch::Approx(170.0));
+}