diff --git a/benchmarks/gemm.cpp b/benchmarks/gemm.cpp index 2086210..b022de9 100644 --- a/benchmarks/gemm.cpp +++ b/benchmarks/gemm.cpp @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 #include +#include #include #include @@ -192,9 +193,9 @@ auto manual_blocked_tiled_subview(const MdSpanA A, const MdSpanB B, MdSpanC C) - assert(A.extent(0) == C.extent(0)); assert(B.extent(1) == C.extent(1)); - using ATile = std::mdspan, std::layout_right>; - using BTile = std::mdspan, std::layout_right>; - using CTile = std::mdspan, std::layout_right>; + using ATile = std::mdspan, typename MdSpanA::mapping_type::inner_layout_type>; + using BTile = std::mdspan, typename MdSpanB::mapping_type::inner_layout_type>; + using CTile = std::mdspan, typename MdSpanC::mapping_type::inner_layout_type>; const size_t M = A.extent(0); const size_t N = B.extent(1); @@ -228,14 +229,14 @@ auto manual_blocked_tiled_subview(const MdSpanA A, const MdSpanB B, MdSpanC C) - template requires IsMdSpan2D && IsMdSpan2D && IsMdSpan2D -auto manual_blocked_tiled_subview_BLTL(const MdSpanA A, const MdSpanB B, MdSpanC C) -> void { +auto manual_blocked_tiled_subview_vec_unroll_k(const MdSpanA A, const MdSpanB B, MdSpanC C) -> void { assert(A.extent(1) == B.extent(0)); assert(A.extent(0) == C.extent(0)); assert(B.extent(1) == C.extent(1)); - using ATile = std::mdspan, std::layout_right>; - using BTile = std::mdspan, std::layout_left>; - using CTile = std::mdspan, std::layout_right>; + using ATile = std::mdspan, typename MdSpanA::mapping_type::inner_layout_type>; + using BTile = std::mdspan, typename MdSpanB::mapping_type::inner_layout_type>; + using CTile = std::mdspan, typename MdSpanC::mapping_type::inner_layout_type>; const size_t M = A.extent(0); const size_t N = B.extent(1); @@ -253,6 +254,52 @@ auto manual_blocked_tiled_subview_BLTL(const MdSpanA A, const MdSpanB B, MdSpanC for (size_t i = 0; i < std::min(TM, M - x); ++i) { for (size_t j = 0; j < std::min(TN, N - y); ++j) { Type acc = 0.0; +#if defined(__AVX__) && defined(__FMA__) + // 8x unrolled and vectorized kernel (processing 32 elements/iter). + // If TO == 32, this effectively fully unrolls the `k` loop. + // + // On Zen3, `vfmaddXXXpd` is 4 lat / 0.5 tp (CPI) => at least 2 FMAs needed to achieve max throughput. + // Unrolling more allows to amortize the cost of the final horizontal reduction. + // + // TODO(dssgabriel): Attempt to vectorize the `j` loop to remove the need for the horizontal reduction. + if constexpr (TO >= 32 and TO % 32 == 0) { + __m256d vacc = _mm256_setzero_pd(); + for (size_t k = 0; k < std::min(TO, O - z); k += 32) { + __m256d va0 = _mm256_loadu_pd(&tA[i, k]); + __m256d vb0 = _mm256_loadu_pd(&tB[k, j]); + vacc = _mm256_fmadd_pd(va0, vb0, vacc); + __m256d va1 = _mm256_loadu_pd(&tA[i, k + 4]); + __m256d vb1 = _mm256_loadu_pd(&tB[k + 4, j]); + vacc = _mm256_fmadd_pd(va1, vb1, vacc); + __m256d va2 = _mm256_loadu_pd(&tA[i, k + 8]); + __m256d vb2 = _mm256_loadu_pd(&tB[k + 8, j]); + vacc = _mm256_fmadd_pd(va2, vb2, vacc); + __m256d va3 = _mm256_loadu_pd(&tA[i, k + 12]); + __m256d vb3 = _mm256_loadu_pd(&tB[k + 12, j]); + vacc = _mm256_fmadd_pd(va3, vb3, vacc); + __m256d va4 = _mm256_loadu_pd(&tA[i, k + 16]); + __m256d vb4 = _mm256_loadu_pd(&tB[k + 16, j]); + vacc = _mm256_fmadd_pd(va4, vb4, vacc); + __m256d va5 = _mm256_loadu_pd(&tA[i, k + 20]); + __m256d vb5 = _mm256_loadu_pd(&tB[k + 20, j]); + vacc = _mm256_fmadd_pd(va5, vb5, vacc); + __m256d va6 = _mm256_loadu_pd(&tA[i, k + 24]); + __m256d vb6 = _mm256_loadu_pd(&tB[k + 24, j]); + vacc = _mm256_fmadd_pd(va6, vb6, vacc); + __m256d va7 = _mm256_loadu_pd(&tA[i, k + 28]); + __m256d vb7 = _mm256_loadu_pd(&tB[k + 28, j]); + vacc = _mm256_fmadd_pd(va7, vb7, vacc); + } + __m128d hsum = _mm_add_pd(_mm256_castpd256_pd128(vacc), _mm256_extractf128_pd(vacc, 1)); + hsum = _mm_hadd_pd(hsum, hsum); + acc = _mm_cvtsd_f64(hsum); + } else { + for (size_t k = 0; k < std::min(TO, O - z); ++k) { + acc += tA[i, k] * tB[k, j]; + } + } + tC[i, j] += acc; +#else #if defined(LAYOUT_TILED_BENCH_ENABLE_OMP) #pragma omp simd #endif @@ -260,6 +307,7 @@ auto manual_blocked_tiled_subview_BLTL(const MdSpanA A, const MdSpanB B, MdSpanC acc += tA[i, k] * tB[k, j]; } tC[i, j] += acc; +#endif } } } @@ -267,6 +315,86 @@ auto manual_blocked_tiled_subview_BLTL(const MdSpanA A, const MdSpanB B, MdSpanC } } +#if defined(__AVX__) && defined(__FMA__) +template + requires IsMdSpan2D && IsMdSpan2D && IsMdSpan2D +auto manual_blocked_tiled_subview_vec_unroll2(const MdSpanA A, const MdSpanB B, MdSpanC C) -> void { + assert(A.extent(1) == B.extent(0)); + assert(A.extent(0) == C.extent(0)); + assert(B.extent(1) == C.extent(1)); + + using ATile = std::mdspan, typename MdSpanA::mapping_type::inner_layout_type>; + using BTile = std::mdspan, typename MdSpanB::mapping_type::inner_layout_type>; + using CTile = std::mdspan, typename MdSpanC::mapping_type::inner_layout_type>; + + constexpr size_t MR = 4; + constexpr size_t NR = 8; + + const size_t M = A.extent(0); + const size_t N = B.extent(1); + const size_t O = A.extent(1); + +#if defined(LAYOUT_TILED_BENCH_ENABLE_OMP) +#pragma omp parallel for schedule(static) +#endif + for (size_t x = 0; x < M; x += TM) { + for (size_t y = 0; y < N; y += TN) { + auto tC = CTile(C.data_handle() + C.mapping().tile_offset_outer(x, y)); + for (size_t z = 0; z < O; z += TO) { + const auto tA = ATile(A.data_handle() + A.mapping().tile_offset_outer(x, z)); + const auto tB = BTile(B.data_handle() + B.mapping().tile_offset_outer(z, y)); + // 4x8 register blocking + for (size_t i = 0; i < TM; i += MR) { + for (size_t j = 0; j < TN; j += NR) { + // 8 accumulators over 4 rows/2 columns + __m256d vc00 = _mm256_loadu_pd(&tC[i, j]); + __m256d vc01 = _mm256_loadu_pd(&tC[i, j + 4]); + __m256d vc10 = _mm256_loadu_pd(&tC[i + 1, j]); + __m256d vc11 = _mm256_loadu_pd(&tC[i + 1, j + 4]); + __m256d vc20 = _mm256_loadu_pd(&tC[i + 2, j]); + __m256d vc21 = _mm256_loadu_pd(&tC[i + 2, j + 4]); + __m256d vc30 = _mm256_loadu_pd(&tC[i + 3, j]); + __m256d vc31 = _mm256_loadu_pd(&tC[i + 3, j + 4]); + + for (size_t k = 0; k < TO; ++k) { + // Broadcast 4 elements of A + __m256d sa0 = _mm256_broadcast_sd(&tA[i, k]); + __m256d sa1 = _mm256_broadcast_sd(&tA[i + 1, k]); + __m256d sa2 = _mm256_broadcast_sd(&tA[i + 2, k]); + __m256d sa3 = _mm256_broadcast_sd(&tA[i + 3, k]); + + // Multiply broadcast A with 2 columns of B + __m256d vb0 = _mm256_loadu_pd(&tB[k, j]); + vc00 = _mm256_fmadd_pd(sa0, vb0, vc00); + vc10 = _mm256_fmadd_pd(sa1, vb0, vc10); + vc20 = _mm256_fmadd_pd(sa2, vb0, vc20); + vc30 = _mm256_fmadd_pd(sa3, vb0, vc30); + __m256d vb1 = _mm256_loadu_pd(&tB[k, j + 4]); + vc01 = _mm256_fmadd_pd(sa0, vb1, vc01); + vc11 = _mm256_fmadd_pd(sa1, vb1, vc11); + vc21 = _mm256_fmadd_pd(sa2, vb1, vc21); + vc31 = _mm256_fmadd_pd(sa3, vb1, vc31); + } + + // Store accumulators + _mm256_storeu_pd(&tC[i, j], vc00); + _mm256_storeu_pd(&tC[i, j + 4], vc01); + _mm256_storeu_pd(&tC[i + 1, j], vc10); + _mm256_storeu_pd(&tC[i + 1, j + 4], vc11); + _mm256_storeu_pd(&tC[i + 2, j], vc20); + _mm256_storeu_pd(&tC[i + 2, j + 4], vc21); + _mm256_storeu_pd(&tC[i + 3, j], vc30); + _mm256_storeu_pd(&tC[i + 3, j + 4], vc31); + } + } + } + } + } +} +#else +#error unsupported impl on target arch +#endif + #if defined(LAYOUT_TILED_BENCH_ENABLE_BLAS) template requires IsMdSpan2D && IsMdSpan2D && IsMdSpan2D && IsLayoutRight @@ -339,10 +467,10 @@ auto run_bench(const Config& cfg) -> int { auto mtlB_l2 = std::mdspan>(tlB_l2.get(), cfg.n_, cfg.n_); // Matrix initializations - ankerl::nanobench::Rng rng_naive(42); - benchmark_utils::fill_2d_uniform01(mrA, rng_naive); - benchmark_utils::fill_2d_uniform01(mrB, rng_naive); - benchmark_utils::fill_2d_uniform01(mlB, rng_naive); + ankerl::nanobench::Rng rng(42); + benchmark_utils::fill_2d_uniform01(mrA, rng); + benchmark_utils::fill_2d_uniform01(mrB, rng); + benchmark_utils::fill_2d_copy(mlB, mrB); benchmark_utils::fill_2d_copy(mtrA_l1, mrA); benchmark_utils::fill_2d_copy(mtrB_l1, mrB); benchmark_utils::fill_2d_copy(mtlB_l1, mlB); @@ -388,6 +516,27 @@ auto run_bench(const Config& cfg) -> int { res = benchmark_utils::check(ref, benchmark_utils::l2_norm(mtrC_l2), "RTR-L2_mbtsv", big_o); benchmark_utils::fill_zeroes(mtrC_l2); + manual_blocked_tiled_subview(mtrA_l1, mtlB_l1, mtrC_l1); + res = benchmark_utils::check(ref, benchmark_utils::l2_norm(mtrC_l1), "RTR-BLTL-L1_mbtsv", big_o); + benchmark_utils::fill_zeroes(mtrC_l1); + manual_blocked_tiled_subview(mtrA_l2, mtlB_l2, mtrC_l2); + res = benchmark_utils::check(ref, benchmark_utils::l2_norm(mtrC_l2), "RTR-BLTL-L2_mbtsv", big_o); + benchmark_utils::fill_zeroes(mtrC_l2); + + manual_blocked_tiled_subview_vec_unroll_k(mtrA_l1, mtlB_l1, mtrC_l1); + res = benchmark_utils::check(ref, benchmark_utils::l2_norm(mtrC_l1), "RTR-BLTL-L1_mbtvu", big_o); + benchmark_utils::fill_zeroes(mtrC_l1); + manual_blocked_tiled_subview_vec_unroll_k(mtrA_l2, mtlB_l2, mtrC_l2); + res = benchmark_utils::check(ref, benchmark_utils::l2_norm(mtrC_l2), "RTR-BLTL-L2_mbtvu", big_o); + benchmark_utils::fill_zeroes(mtrC_l2); + + manual_blocked_tiled_subview_vec_unroll2(mtrA_l1, mtrB_l1, mtrC_l1); + res = benchmark_utils::check(ref, benchmark_utils::l2_norm(mtrC_l1), "RTR-BLTL-L1_mbt_vu2", big_o); + benchmark_utils::fill_zeroes(mtrC_l1); + manual_blocked_tiled_subview_vec_unroll2(mtrA_l2, mtrB_l2, mtrC_l2); + res = benchmark_utils::check(ref, benchmark_utils::l2_norm(mtrC_l2), "RTR-BLTL-L2_mbt_vu2", big_o); + benchmark_utils::fill_zeroes(mtrC_l2); + return res; #else // Benchmarks @@ -438,10 +587,22 @@ auto run_bench(const Config& cfg) -> int { manual_blocked_tiled_subview(mtrA_l2, mtrB_l2, mtrC_l2); }); bs.run(fmt::format("RTR_BLTL-{0}x{0} mbtsv", TILE_L1), [&]() { - manual_blocked_tiled_subview_BLTL(mtrA_l1, mtlB_l1, mtrC_l1); + manual_blocked_tiled_subview(mtrA_l1, mtlB_l1, mtrC_l1); }); bs.run(fmt::format("RTR_BLTL-{0}x{0} mbtsv", TILE_L2), [&]() { - manual_blocked_tiled_subview_BLTL(mtrA_l2, mtlB_l2, mtrC_l2); + manual_blocked_tiled_subview(mtrA_l2, mtlB_l2, mtrC_l2); + }); + bs.run(fmt::format("RTR_BLTL-{0}x{0} mbt_vu", TILE_L1), [&]() { + manual_blocked_tiled_subview_vec_unroll_k(mtrA_l1, mtlB_l1, mtrC_l1); + }); + bs.run(fmt::format("RTR_BLTL-{0}x{0} mbt_vu", TILE_L2), [&]() { + manual_blocked_tiled_subview_vec_unroll_k(mtrA_l2, mtlB_l2, mtrC_l2); + }); + bs.run(fmt::format("RTR_BLTL-{0}x{0} mbt_vu2", TILE_L1), [&]() { + manual_blocked_tiled_subview_vec_unroll2(mtrA_l1, mtrB_l1, mtrC_l1); + }); + bs.run(fmt::format("RTR_BLTL-{0}x{0} mbt_vu2", TILE_L2), [&]() { + manual_blocked_tiled_subview_vec_unroll2(mtrA_l2, mtrB_l2, mtrC_l2); }); // Output