From 5ac9489c079f3f1a0ba1d2d8385001ada83a13b8 Mon Sep 17 00:00:00 2001 From: y-jan137 Date: Fri, 13 Mar 2026 16:26:47 +0300 Subject: Implement vectorized matmul --- src/matrix.cpp | 152 +++++++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 147 insertions(+), 5 deletions(-) (limited to 'src/matrix.cpp') 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 +#include #include #include +#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && \ + (defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)) +#include +#endif + +#if !defined(LINEAR_ALGEBRA_FORCE_SCALAR_MATMUL) && defined(__ARM_NEON) && defined(__aarch64__) && !defined(__clangd__) +#include +#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; -- cgit v1.2.3