From 428e2ce6986df59a4e202d85542a1dd7fbb5fd32 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Wed, 29 Jul 2026 17:43:07 -0700 Subject: [PATCH 01/12] [blas] Vectorize bf16-in/float-out gemm for custom SDPA custom SDPA's q@k.T and attn@v matmuls (bf16 inputs, fp32 accumulation) instantiated the generic scalar gemm_{transa,notrans}_ templates -- no vectorization, so bf16 SDPA ran ~5.5x slower than fp32. Add optimized specializations that use the existing vectorized internal::bf16_dot_with_fp32_arith (fp32 accumulate, identical numerics to the scalar path). notrans packs a's strided-k rows into a contiguous buffer so the dot applies. Kept serial: custom SDPA already parallelizes its outer head loop and executorch's threadpool does not support nested parallelism. gemma-3-1b bf16 8da4w: custom_sdpa 668->372ms, Method::execute 977->669ms (prefill @16c 974->670ms); argmax unchanged (107). --- kernels/optimized/blas/BlasKernel.h | 63 +++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/kernels/optimized/blas/BlasKernel.h b/kernels/optimized/blas/BlasKernel.h index 9d4a9f87ec0..67bcf99b909 100644 --- a/kernels/optimized/blas/BlasKernel.h +++ b/kernels/optimized/blas/BlasKernel.h @@ -15,6 +15,7 @@ #include #include +#include namespace executorch { namespace cpublas { @@ -140,6 +141,42 @@ float bf16_dot_with_fp32_arith( int64_t len); } // namespace internal +// Used by custom SDPA's attn@V. Serial on purpose: SDPA already parallelizes +// its outer head loop, and executorch's threadpool deadlocks on a nested +// parallel_for. +// clang-format off +template <> +inline typename std::enable_if< + !std::is_same::value, + void>::type +gemm_notrans_( + int64_t m, int64_t n, int64_t k, + float alpha, + const torch::executor::BFloat16 *a, int64_t lda, + const torch::executor::BFloat16 *b, int64_t ldb, + float beta, + float *c, int64_t ldc) { + // a's k-dimension is strided by lda; bf16_dot_with_fp32_arith needs it + // contiguous. + std::vector a_row(k); + for (int64_t i = 0; i < m; ++i) { + for (int64_t l = 0; l < k; ++l) { + a_row[l] = a[l * lda + i]; + } + const torch::executor::BFloat16 *b_ = b; + for (int64_t j = 0; j < n; ++j) { + const float dot = internal::bf16_dot_with_fp32_arith(a_row.data(), b_, k); + b_ += ldb; + if (beta == 0) { + c[j * ldc + i] = alpha * dot; + } else { + c[j * ldc + i] = beta * c[j * ldc + i] + alpha * dot; + } + } + } +} +// clang-format on + // clang-format off template void gemm_transa_( @@ -209,6 +246,32 @@ inline void gemm_transa_ +inline void gemm_transa_( + int64_t m, int64_t n, int64_t k, + float alpha, + const torch::executor::BFloat16 *a, int64_t lda, + const torch::executor::BFloat16 *b, int64_t ldb, + float beta, + float *c, int64_t ldc) { + const auto *a_ = a; + for (int i = 0; i < m; ++i) { + const auto *b_ = b; + for (int j = 0; j < n; ++j) { + const float dot = internal::bf16_dot_with_fp32_arith(a_, b_, k); + b_ += ldb; + if (beta == 0) { + c[j*ldc+i] = alpha*dot; + } else { + c[j*ldc+i] = beta*c[j*ldc+i]+alpha*dot; + } + } + a_ += lda; + } +} // clang-format on template From a39681b0ea5aea71cf213dca3573897dd31f4b65 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 10 Aug 2026 16:20:27 -0400 Subject: [PATCH 02/12] [blas] Initialize cpuinfo before querying it for the bf16 dot cpuinfo_has_* reads a zeroed struct until cpuinfo_initialize() has run, so bf16_dot_with_fp32_arith silently took the widening fallback for any caller that had not already initialized cpuinfo elsewhere -- the threadpool, say. Nothing failed; the fast dot just never ran. This is not an x86 concern despite arriving alongside one: the check it fixes is the aarch64 one, and an unnoticed fallback there costs the whole bfdot path. Split out of the AVX512-BF16 commit so architectures that do not want the x86 dot still get it. Authored with Claude Code. --- kernels/optimized/blas/BlasKernel.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/kernels/optimized/blas/BlasKernel.cpp b/kernels/optimized/blas/BlasKernel.cpp index 970105c859e..5730d114a08 100644 --- a/kernels/optimized/blas/BlasKernel.cpp +++ b/kernels/optimized/blas/BlasKernel.cpp @@ -333,7 +333,8 @@ float bf16_dot_with_fp32_arith( const at::BFloat16* vec2, int64_t len) { #if COMPILER_SUPPORTS_BF16_TARGET - if (cpuinfo_has_arm_bf16()) { + // cpuinfo_has_* reads a zeroed struct until cpuinfo_initialize() runs. + if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { return dot_with_fp32_arith_bfdot(vec1, vec2, len); } else #endif // COMPILER_SUPPORTS_BF16_TARGET From 53aedaf446db45c343ea6effdab3b8de6aebb590 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 3 Aug 2026 09:46:29 -0700 Subject: [PATCH 03/12] [blas] Tests for the bf16 dot and bf16-in/float-out gemm paths Cover internal::bf16_dot_with_fp32_arith and the gemm_{transa,notrans}_ specializations against a sequential fp32 reference. Lengths straddle the vector-loop, cleanup-loop and scalar-tail boundaries of both bfdot paths (128/32 bf16 per iteration on x86, 32/8 on ARM), the gemm cases use padded leading dimensions and a range of alpha/beta, and a separate case pins the beta == 0 overwrite semantics with a NaN-filled output. Verified the tests bite: with the x86 scalar tail loop deleted they fail at exactly the non-multiple-of-32 lengths, and they pass with dispatch forced to the portable fallback. That mutation is also what caught the missing cpuinfo_initialize() fixed in the previous commit -- before it, the tests passed no matter what the AVX512 path computed. Only the dot implementation the host dispatches to is covered by a given run. --- kernels/optimized/test/libblas_test.cpp | 185 ++++++++++++++++++++++++ 1 file changed, 185 insertions(+) diff --git a/kernels/optimized/test/libblas_test.cpp b/kernels/optimized/test/libblas_test.cpp index 3c944214417..358db95a672 100644 --- a/kernels/optimized/test/libblas_test.cpp +++ b/kernels/optimized/test/libblas_test.cpp @@ -11,6 +11,10 @@ #include #include +#include +#include +#include +#include #include #define TEST_FORALL_SUPPORTED_CTYPES(_, N) \ @@ -45,6 +49,68 @@ bool check_all_equal_to(std::vector& arr, const float value) { return true; } +template +std::vector make_values(size_t n, uint32_t seed) { + std::vector v(n); + uint32_t state = seed; + for (size_t i = 0; i < n; ++i) { + state = state * 1664525u + 1013904223u; + v[i] = static_cast( + static_cast(state >> 8) / static_cast(1 << 23) - 1.0f); + } + return v; +} + +template +float reference_dot(const T* a, const T* b, int64_t len) { + float sum = 0; + for (int64_t i = 0; i < len; ++i) { + sum += static_cast(a[i]) * static_cast(b[i]); + } + return sum; +} + +// Column-major c = beta * c + alpha * (op(a) @ b), accumulated in fp32, mirror- +// ing the generic gemm_{transa,notrans}_ templates the bf16 specializations +// replace. +template +void reference_gemm( + bool transa, + int64_t m, + int64_t n, + int64_t k, + float alpha, + const T* a, + int64_t lda, + const T* b, + int64_t ldb, + float beta, + std::vector& c, + int64_t ldc) { + for (int64_t i = 0; i < m; ++i) { + for (int64_t j = 0; j < n; ++j) { + float dot = 0; + for (int64_t l = 0; l < k; ++l) { + const float av = transa ? static_cast(a[i * lda + l]) + : static_cast(a[l * lda + i]); + dot += av * static_cast(b[j * ldb + l]); + } + c[j * ldc + i] = + beta == 0 ? alpha * dot : beta * c[j * ldc + i] + alpha * dot; + } + } +} + +void expect_near_relative(float actual, float expected, const char* context) { + EXPECT_NEAR(actual, expected, 1e-4f * std::max(1.0f, std::abs(expected))) + << context; +} + +// Straddle the vectorized main-loop, cleanup-loop and scalar-tail boundaries of +// the bfdot paths: 128 and 32 bf16 per iteration on x86, 32 and 8 on ARM. +constexpr int64_t kDotLengths[] = + {0, 1, 7, 8, 15, 16, 31, 32, 33, 63, 64, 96, 127, 128, 129, 160, 255, 257}; + } // namespace template @@ -80,3 +146,122 @@ void test_matmul_ones() { TEST(BlasTest, MatmulOnes) { TEST_FORALL_SUPPORTED_CTYPES(test_matmul_ones, 25); } + +// bf16_dot_with_fp32_arith has three implementations -- ARM bfdot, x86 +// AVX512-BF16 and the portable fp32 fallback -- selected by compile-time +// support and runtime cpuinfo. Only the one this host dispatches to is covered +// by a given run. +TEST(BlasTest, BF16DotMatchesScalarAccumulation) { + using torch::executor::BFloat16; + + for (const int64_t len : kDotLengths) { + const auto a = make_values(len, 1); + const auto b = make_values(len, 2); + + const float actual = + executorch::cpublas::internal::bf16_dot_with_fp32_arith( + a.data(), b.data(), len); + + expect_near_relative( + actual, + reference_dot(a.data(), b.data(), len), + ("len=" + std::to_string(len)).c_str()); + } +} + +// The bf16-in/float-out gemm specializations used by custom SDPA. +TEST(BlasTest, BF16FloatGemmMatchesScalarAccumulation) { + using executorch::aten::BFloat16; + using executorch::cpublas::TransposeType; + + constexpr int64_t kM = 3; + constexpr int64_t kN = 5; + + for (const bool transa : {false, true}) { + for (const int64_t k : kDotLengths) { + if (k == 0) { + continue; + } + // Pad the leading dimensions so a stride bug can't hide behind tightly + // packed operands. + const int64_t lda = (transa ? k : kM) + 2; + const int64_t ldb = k + 3; + const int64_t ldc = kM + 1; + + const auto a = make_values(lda * (transa ? kM : k), 3); + const auto b = make_values(ldb * kN, 4); + + for (const float alpha : {1.0f, -0.5f}) { + for (const float beta : {0.0f, 1.0f, 0.25f}) { + auto c = make_values(ldc * kN, 5); + auto expected = c; + + // clang-format off + reference_gemm( + transa, + kM, kN, k, + alpha, + a.data(), lda, + b.data(), ldb, + beta, + expected, ldc); + + executorch::cpublas::gemm( + transa ? TransposeType::Transpose : TransposeType::NoTranspose, + TransposeType::NoTranspose, + kM, kN, k, + alpha, + a.data(), lda, + b.data(), ldb, + beta, + c.data(), ldc); + // clang-format on + + const std::string context = "transa=" + std::to_string(transa) + + " k=" + std::to_string(k) + " alpha=" + std::to_string(alpha) + + " beta=" + std::to_string(beta); + for (int64_t j = 0; j < kN; ++j) { + for (int64_t i = 0; i < kM; ++i) { + expect_near_relative( + c[j * ldc + i], expected[j * ldc + i], context.c_str()); + } + } + } + } + } + } +} + +// beta == 0 must overwrite c rather than read it, so uninitialized garbage in +// the output cannot poison the result. +TEST(BlasTest, BF16FloatGemmBetaZeroIgnoresOutput) { + using executorch::aten::BFloat16; + using executorch::cpublas::TransposeType; + + constexpr int64_t kM = 3; + constexpr int64_t kN = 5; + constexpr int64_t kK = 40; + + const auto a = make_values(kK * kM, 6); + const auto b = make_values(kK * kN, 7); + + for (const bool transa : {false, true}) { + std::vector c(kM * kN, std::numeric_limits::quiet_NaN()); + + // clang-format off + executorch::cpublas::gemm( + transa ? TransposeType::Transpose : TransposeType::NoTranspose, + TransposeType::NoTranspose, + kM, kN, kK, + 1.0f, + a.data(), transa ? kK : kM, + b.data(), kK, + 0.0f, + c.data(), kM); + // clang-format on + + for (const float v : c) { + EXPECT_FALSE(std::isnan(v)) << "transa=" << transa; + } + } +} From 0f3c242c66f6c2564bd6a63704edb145879b4872 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 10 Aug 2026 16:00:22 -0400 Subject: [PATCH 04/12] [blas] Skip the row gather in bf16 attn@V when n == 1 gemm_notrans_ gathers each of a's rows into a contiguous buffer so bf16_dot_with_fp32_arith can consume it. The gather costs k strided loads and buys n dots, so at n == 1 it is one gather per multiply-add. Custom SDPA hits exactly that at decode, where the q block is a single row. a's m-dimension is already contiguous, so the alternative needs no gather: convert to fp32 and accumulate along it, with l outermost so a's column is loaded once and reused across c's columns. Per-tile at the decode shape, the gather-free form is 29x faster on Genoa (AVX512-BF16) and roughly 4x on an M4 Max. Above n == 1 the gather amortizes and is kept, unchanged. Authored with Claude Code. --- kernels/optimized/blas/BlasKernel.h | 46 +++++++++++++++++++++++++++-- 1 file changed, 44 insertions(+), 2 deletions(-) diff --git a/kernels/optimized/blas/BlasKernel.h b/kernels/optimized/blas/BlasKernel.h index 67bcf99b909..7acb503f9de 100644 --- a/kernels/optimized/blas/BlasKernel.h +++ b/kernels/optimized/blas/BlasKernel.h @@ -156,8 +156,50 @@ gemm_notrans_( const torch::executor::BFloat16 *b, int64_t ldb, float beta, float *c, int64_t ldc) { - // a's k-dimension is strided by lda; bf16_dot_with_fp32_arith needs it - // contiguous. + // bf16_dot_with_fp32_arith needs a contiguous k-vector, but a is strided by + // lda in k, so the dot path must gather each row first: k strided loads to + // buy n dots. At n == 1 that is one gather per multiply-add and cannot pay + // on any architecture. a's m-dimension is contiguous, so accumulate along it + // instead, which is the traversal the fp32 specialization already uses. + constexpr int64_t kMinColsForGather = 2; + const bool use_dot = n >= kMinColsForGather; + + if (!use_dot) { + for (int64_t j = 0; j < n; ++j) { + float *c_col = c + j * ldc; + if (beta == 0) { + for (int64_t i = 0; i < m; ++i) { + c_col[i] = 0.0f; + } + } else if (beta != 1.0f) { + for (int64_t i = 0; i < m; ++i) { + c_col[i] *= beta; + } + } + } + // l outermost so a's column is loaded once and reused across c's columns. + for (int64_t l = 0; l < k; ++l) { + const torch::executor::BFloat16 *a_col = a + l * lda; + for (int64_t j = 0; j < n; ++j) { + const float b_val = static_cast(b[l + j * ldb]) * alpha; + float *c_col = c + j * ldc; + // Unrolled for the same reason the fp32 specialization above is: the + // bf16->fp32 conversion does not vectorize from the rolled form. + int64_t i = 0; + for (; i + 4 <= m; i += 4) { + c_col[i + 0] += static_cast(a_col[i + 0]) * b_val; + c_col[i + 1] += static_cast(a_col[i + 1]) * b_val; + c_col[i + 2] += static_cast(a_col[i + 2]) * b_val; + c_col[i + 3] += static_cast(a_col[i + 3]) * b_val; + } + for (; i < m; ++i) { + c_col[i] += static_cast(a_col[i]) * b_val; + } + } + } + return; + } + std::vector a_row(k); for (int64_t i = 0; i < m; ++i) { for (int64_t l = 0; l < k; ++l) { From 03be6cf0f9e389418bb56dd5e7a5f3fa869a8b43 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 10 Aug 2026 15:57:04 -0400 Subject: [PATCH 05/12] [blas] Widen bf16 q@K.T to fp32 and let BLAS run it BLAS has no bf16 entry point, so a bf16 gemm falls to our own kernels while the fp32 overload gets a packed, blocked BLAS. Custom SDPA calls both, so q@K.T -- a short reduction (headSize) feeding an m*n tile -- was a per-output dot competing against a packed BLAS, paying a cross-lane reduction per result. Widen the operands and use the same BLAS. Conversion is O(mk + nk) against O(mnk) of multiply, so it only pays for a short k with enough output columns; the thresholds are empirical and decode's single column is excluded. Done at the SDPA call site rather than inside cpublas so the scratch comes from ctx.allocate_temp alongside the buffers cpu_flash_attention already allocates, rather than a thread_local in shared kernel code that memory planning cannot see. cpublas exposes gemm_uses_blas() so the caller can tell whether widening is worth it; the build flag that answers that is only visible inside that library. BFloat16 only: the widened path reinterprets the operands, and Half was never measured through it. Authored with Claude Code. --- extension/llm/custom_ops/op_sdpa_impl.h | 72 ++++++++++++++++++++++++- kernels/optimized/blas/CPUBlas.cpp | 8 +++ kernels/optimized/blas/CPUBlas.h | 7 +++ 3 files changed, 85 insertions(+), 2 deletions(-) diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index 9c479569505..2c6fb09ed49 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -71,7 +71,8 @@ void _q_at_k_gemm( const int64_t q_stride_m, const MaybeQuantizedMatrixData& k_data, const int64_t k_stride_n, - accum_t* qk_data) { + accum_t* qk_data, + accum_t* widen_scratch) { ET_CHECK_MSG(q_data.dtype == k_data.dtype, "q and k must have same dtype"); ET_CHECK_MSG( q_data.dtype == ScalarType::Char || q_data.dtype == ScalarType::Float || @@ -109,6 +110,44 @@ void _q_at_k_gemm( q_data.dtype == ScalarType::BFloat16 || q_data.dtype == ScalarType::Half) { if constexpr (std::is_same::value) { + if (widen_scratch != nullptr && + q_data.dtype == ScalarType::BFloat16) { + // Reduction is headSize here, short enough that a per-output dot pays + // a cross-lane reduction it cannot amortize, while a packed BLAS can. + // BLAS has no bf16 entry point, so widen and use the fp32 one; the + // conversion is O(mk + nk) against O(mnk) of multiply. + accum_t* k_f32 = widen_scratch; + accum_t* q_f32 = widen_scratch + k_n * qk_k; + const auto* k_src = + static_cast(k_data.data); + const auto* q_src = + static_cast(q_data.data); + for (int64_t i = 0; i < k_n; ++i) { + for (int64_t l = 0; l < qk_k; ++l) { + k_f32[i * qk_k + l] = static_cast(k_src[i * k_stride_n + l]); + } + } + for (int64_t j = 0; j < q_m; ++j) { + for (int64_t l = 0; l < qk_k; ++l) { + q_f32[j * qk_k + l] = static_cast(q_src[j * q_stride_m + l]); + } + } + ::executorch::cpublas::gemm( + ::executorch::cpublas::TransposeType::Transpose, + ::executorch::cpublas::TransposeType::NoTranspose, + k_n, + q_m, + qk_k, + static_cast(1), + k_f32, + qk_k, + q_f32, + qk_k, + static_cast(0), + qk_data, + k_n); + return; + } auto do_gemm = [&](auto rt_tag) { using rt = decltype(rt_tag); ::executorch::cpublas::gemm( @@ -848,6 +887,29 @@ void cpu_flash_attention( buf_reduced = scratch_reduced.get(); } } + // Scratch for widening q@K.T to fp32 (see _q_at_k_gemm): one K block plus one + // q block. BFloat16 only -- the widened path reinterprets the operands. + const bool widen_qk = + std::is_same::value && + ::executorch::cpublas::gemm_uses_blas(); + // Widening only pays for a short reduction feeding many output columns; both + // bounds are empirical. + constexpr int64_t kMaxHeadSizeForWidenedQK = 128; + constexpr int64_t kMinQBlockForWidenedQK = 32; + int64_t size_per_thread_widen = widen_qk ? (kvSplitSize + qSplitSize) * headSize : 0; + std::unique_ptr allocated_buf_widen; + accum_t* widen_buf = nullptr; + if (widen_qk) { + int64_t size_widen_bytes = size_per_thread_widen * num_thread * sizeof(accum_t); + Result scratch_widen = ctx.allocate_temp(size_widen_bytes, 64); + if (!scratch_widen.ok()) { + allocated_buf_widen = std::make_unique(size_widen_bytes); + widen_buf = reinterpret_cast(allocated_buf_widen.get()); + } else { + widen_buf = reinterpret_cast(scratch_widen.get()); + } + } + int64_t size_per_thread_qdq_vec = kvSplitSize * headSize; // Lets align size_per_thread_qdq_vec to 64 bytes, for coalesced cache reads, // by padding with right number of per thread elements @@ -891,6 +953,8 @@ void cpu_flash_attention( : nullptr; accum_t* buf_qdq_ptr = scratch_for_quant_dequant + ompIdx * size_per_thread_qdq_vec; + accum_t* widen_ptr = + widen_buf ? widen_buf + ompIdx * size_per_thread_widen : nullptr; for (int64_t z = begin; z < end; z++) { int64_t m = k * qSplitSize; @@ -987,7 +1051,11 @@ void cpu_flash_attention( qStrideM, k_sub_matrix_data, kStrideN, - qk_data); + qk_data, + (widen_qk && headSize <= kMaxHeadSizeForWidenedQK && + qBlockSize >= kMinQBlockForWidenedQK) + ? widen_ptr + : nullptr); // There are 4 cases that is_causal has to cover to fill // not-attendable-position with -inf diff --git a/kernels/optimized/blas/CPUBlas.cpp b/kernels/optimized/blas/CPUBlas.cpp index 64270b84a20..32e0b2c3a86 100644 --- a/kernels/optimized/blas/CPUBlas.cpp +++ b/kernels/optimized/blas/CPUBlas.cpp @@ -47,6 +47,14 @@ inline CBLAS_TRANSPOSE to_cblas_transpose(TransposeType trans) { #endif // ET_BUILD_FOR_APPLE #endif // ET_BUILD_WITH_BLAS +bool gemm_uses_blas() { +#ifdef ET_BUILD_WITH_BLAS + return true; +#else + return false; +#endif +} + // clang-format off void normalize_last_dims( TransposeType transa, TransposeType transb, diff --git a/kernels/optimized/blas/CPUBlas.h b/kernels/optimized/blas/CPUBlas.h index 30e022d3315..8384b72e228 100644 --- a/kernels/optimized/blas/CPUBlas.h +++ b/kernels/optimized/blas/CPUBlas.h @@ -43,6 +43,13 @@ inline char to_blas(TransposeType trans) { return 'N'; } +// Whether gemm() dispatches to a BLAS implementation rather than the portable +// fallback. Callers that can restructure a shape BLAS handles better -- by +// widening a reduced-precision operand, say -- need this to know whether that +// is worth doing. Answered here because the build flag is only visible inside +// this library. +bool gemm_uses_blas(); + // clang-format off template void gemm_impl( From c1f22773aa82c5a57f5d81c3037759243db1fa8f Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 10 Aug 2026 16:00:34 -0400 Subject: [PATCH 06/12] [blas] Raise the bf16 attn@V gather threshold on aarch64 On aarch64 the gather-free form keeps winning well past n == 1. Measured on an M4 Max (m=64, k=512, GFLOP/s, runtime-valued dimensions): the gathered dot runs 6 at n=1, 28 at n=8, 37 at n=32 and 47 at n=256, against a fairly flat 25-42 without the gather, so the crossover sits near n=32. x86 keeps the previous behaviour. vdpbf16ps is much stronger there relative to fma, a per-tile measurement on Genoa put the gathered dot ahead by 1.19x at n=256, and this host cannot measure that case. Note the dimensions must be runtime values to see this: with them constant-folded, clang unrolls the gather-free inner loop to ~60 GFLOP/s and the crossover disappears, which is not what the kernel sees. Authored with Claude Code. --- kernels/optimized/blas/BlasKernel.h | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/kernels/optimized/blas/BlasKernel.h b/kernels/optimized/blas/BlasKernel.h index 7acb503f9de..5b779d1b764 100644 --- a/kernels/optimized/blas/BlasKernel.h +++ b/kernels/optimized/blas/BlasKernel.h @@ -161,7 +161,16 @@ gemm_notrans_( // buy n dots. At n == 1 that is one gather per multiply-add and cannot pay // on any architecture. a's m-dimension is contiguous, so accumulate along it // instead, which is the traversal the fp32 specialization already uses. + // + // On aarch64 the gather-free form keeps winning well past n == 1. On x86 + // vdpbf16ps makes the dot strong enough that only n == 1 avoids it. The + // crossover is empirical, and only visible with runtime-valued dimensions: + // constant-folded ones let the gather-free loop unroll and hide it. +#if defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) + constexpr int64_t kMinColsForGather = 32; +#else constexpr int64_t kMinColsForGather = 2; +#endif const bool use_dot = n >= kMinColsForGather; if (!use_dot) { From 95eb4c85653817ecfe64f65afc7ff1d59b125497 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 4 Aug 2026 11:43:06 -0400 Subject: [PATCH 07/12] [blas] Block bf16 q@K.T by four output columns gemm_transa_ issued one bf16_dot_with_fp32_arith per output element. Custom SDPA calls it with k = headSize -- 64 for a typical LLM -- while producing an m x n tile of 512 x 256, so each of the 131k dots per tile pays its own cross-lane reduction over a reduction length of only 64. At that length the horizontal add is a large fraction of the dot, and it is paid once per output. Add bf16_dot4_with_fp32_arith: four dots against a shared vec1, accumulating into four registers that collapse with a single vpaddq tree rather than four independent cross-lane reductions. gemm_transa_ consumes four columns at a time and falls back to the single dot for the remainder. Measured on an M4 Max at the shape and strides SDPA uses (m=512, n=256, k=64, lda=512, ldb=2048, runtime-valued dimensions): 22.9 -> 42.6 GFLOP/s, against 9.8 for the fp32 specialization. End to end on Llama 3.2 1B 8da4w at 2048 context, bf16 relative to fp32: decode 1.09x -> 1.15-1.19x, prefill 0.66x -> 0.70x, and TTFT from 1.51x to 1.42x worse. Only ratios are quoted because macOS cannot pin cores and the absolute numbers drifted 13% between runs; the ratios reproduced across two. Only the aarch64 bfdot path is blocked. Elsewhere the new entry point issues four sequential dots, which is exactly what the caller did before, so nothing changes there. The same reduction sharing should apply to vdpbf16ps, but there is no AVX512-BF16 host here to measure it on. libblas_test passes. Authored with Claude Code. --- kernels/optimized/blas/BlasKernel.cpp | 56 +++++++++++++++++++++++++++ kernels/optimized/blas/BlasKernel.h | 32 ++++++++++----- 2 files changed, 79 insertions(+), 9 deletions(-) diff --git a/kernels/optimized/blas/BlasKernel.cpp b/kernels/optimized/blas/BlasKernel.cpp index 5730d114a08..153493e32b1 100644 --- a/kernels/optimized/blas/BlasKernel.cpp +++ b/kernels/optimized/blas/BlasKernel.cpp @@ -343,4 +343,60 @@ float bf16_dot_with_fp32_arith( } } +#if COMPILER_SUPPORTS_BF16_TARGET +// Four dots against a shared vec1: the cross-lane reduction is paid once per +// four results rather than per result, which matters when callers reduce over +// headSize while producing a whole tile of outputs. +TARGET_ARM_BF16_ATTRIBUTE static void dot4_with_fp32_arith_bfdot( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t stride2, + int64_t len, + float* out) { + constexpr int64_t kElementsPerIteration = 8; + float32x4_t acc[4] = { + vdupq_n_f32(0.0f), + vdupq_n_f32(0.0f), + vdupq_n_f32(0.0f), + vdupq_n_f32(0.0f)}; + int64_t idx = 0; + for (; idx + kElementsPerIteration <= len; idx += kElementsPerIteration) { + // See NOTE[Intrinsics in bfdot variant] above. + const auto v1 = + vld1q_bf16(reinterpret_cast(&vec1[idx])); + for (int64_t j = 0; j < 4; ++j) { + const auto v2 = vld1q_bf16( + reinterpret_cast(&vec2[j * stride2 + idx])); + acc[j] = vbfdotq_f32(acc[j], v1, v2); + } + } + const float32x4_t sums = vpaddq_f32( + vpaddq_f32(acc[0], acc[1]), vpaddq_f32(acc[2], acc[3])); + vst1q_f32(out, sums); + for (; idx < len; ++idx) { + for (int64_t j = 0; j < 4; ++j) { + out[j] += static_cast(vec1[idx]) * + static_cast(vec2[j * stride2 + idx]); + } + } +} +#endif // COMPILER_SUPPORTS_BF16_TARGET + +void bf16_dot4_with_fp32_arith( + const at::BFloat16* vec1, + const at::BFloat16* vec2, + int64_t stride2, + int64_t len, + float* out) { +#if COMPILER_SUPPORTS_BF16_TARGET + if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { + dot4_with_fp32_arith_bfdot(vec1, vec2, stride2, len, out); + return; + } +#endif // COMPILER_SUPPORTS_BF16_TARGET + for (int64_t j = 0; j < 4; ++j) { + out[j] = bf16_dot_with_fp32_arith(vec1, vec2 + j * stride2, len); + } +} + } // namespace executorch::cpublas::internal diff --git a/kernels/optimized/blas/BlasKernel.h b/kernels/optimized/blas/BlasKernel.h index 5b779d1b764..850bc72049e 100644 --- a/kernels/optimized/blas/BlasKernel.h +++ b/kernels/optimized/blas/BlasKernel.h @@ -139,6 +139,13 @@ float bf16_dot_with_fp32_arith( const torch::executor::BFloat16* vec1, const torch::executor::BFloat16* vec2, int64_t len); +// Four dots of vec1 against vec2, vec2 + stride2, ... into out[0..3]. +void bf16_dot4_with_fp32_arith( + const torch::executor::BFloat16* vec1, + const torch::executor::BFloat16* vec2, + int64_t stride2, + int64_t len, + float* out); } // namespace internal // Used by custom SDPA's attn@V. Serial on purpose: SDPA already parallelizes @@ -308,18 +315,25 @@ inline void gemm_transa_( const torch::executor::BFloat16 *b, int64_t ldb, float beta, float *c, int64_t ldc) { + // Four columns at a time: k is headSize here, short enough that each dot's + // cross-lane reduction is a large share of its cost, and this tile produces + // m*n of them. const auto *a_ = a; - for (int i = 0; i < m; ++i) { - const auto *b_ = b; - for (int j = 0; j < n; ++j) { - const float dot = internal::bf16_dot_with_fp32_arith(a_, b_, k); - b_ += ldb; - if (beta == 0) { - c[j*ldc+i] = alpha*dot; - } else { - c[j*ldc+i] = beta*c[j*ldc+i]+alpha*dot; + for (int64_t i = 0; i < m; ++i) { + int64_t j = 0; + for (; j + 4 <= n; j += 4) { + float dots[4]; + internal::bf16_dot4_with_fp32_arith(a_, b + j * ldb, ldb, k, dots); + for (int64_t d = 0; d < 4; ++d) { + float *dst = c + (j + d) * ldc + i; + *dst = (beta == 0) ? alpha * dots[d] : beta * *dst + alpha * dots[d]; } } + for (; j < n; ++j) { + const float dot = internal::bf16_dot_with_fp32_arith(a_, b + j * ldb, k); + float *dst = c + j * ldc + i; + *dst = (beta == 0) ? alpha * dot : beta * *dst + alpha * dot; + } a_ += lda; } } From 35f06bcd7a87c7328fd9e9a48eeb54da84ffb0cd Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 10 Aug 2026 15:57:19 -0400 Subject: [PATCH 08/12] [llm] Keep attention weights in fp32 and widen V for large q blocks For a reduced-precision activation dtype, cpu_flash_attention casts the fp32 attention weights down to the activation dtype so attn@V can run as a reduced-in/float-out gemm. That cast only earns its keep if the bf16 attn@V kernel is the one that runs. Above a q block of kMinQBlockForWidenedAV it is faster to leave the weights alone, widen V instead, and let BLAS do the multiply -- which also deletes the cast, so the two effects compound. It is also more accurate, because the weights are no longer rounded to bf16. On a 1984-token prefill the logits move measurably toward the fp32 reference (mean |logit| 1.596078 native, 1.591862 widened, against fp32's 1.591681), argmax unchanged. That is one sample, not a numerics validation. Threshold 64, from the q block sizes that actually occur (qSplitSize is 32, 64 or 256 by sequence length, and 1 at decode). Back-to-back A/B on one build, prefill: 256 gains 7.5-8% (2891 -> 2676ms, 2960 -> 2725ms), 64 gains 2-5%, 32 is unresolvable either way, and decode's single output column costs 15% -- the same failure mode as every other conversion that cannot amortize. Gating at 64 keeps decode at 1.19-1.31x of fp32 while prefill reaches parity. Scratch is buf_qdq_ptr rather than a new allocation: it is already per-thread, ctx.allocate_temp'd, 64-byte aligned and sized kvSplitSize * headSize, which is exactly what a widened V needs, and only the quantized branch uses it. Restricted to BFloat16. Half takes the existing path -- the same argument probably applies, but it was not measured. Authored with Claude Code. --- extension/llm/custom_ops/op_sdpa_impl.h | 44 +++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index 2c6fb09ed49..72fd8a60008 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -328,7 +328,8 @@ void _qk_at_v_gemm( accum_t* o_data, const int64_t o_stride_m, const accum_t beta, - accum_t* buf_qdq_ptr) { + accum_t* buf_qdq_ptr, + const bool widen_v) { if (v_data.dtype == ScalarType::Char) { if constexpr (std::is_same::value) { const float* qk = static_cast(qk_data); @@ -378,6 +379,34 @@ void _qk_at_v_gemm( // qk has been cast to the activation dtype (see qk_reduced_data); both // operands are reduced precision and accumulate into the float output. if constexpr (std::is_same::value) { + if (widen_v) { + // Weights are still accum_t, so widen V to match and let BLAS do the + // multiply. buf_qdq_ptr is free here: per-thread, ctx-allocated, the + // right size, and only the quantized branch above uses it. + const auto* v_src = + static_cast(v_data.data); + for (int64_t kk = 0; kk < k; ++kk) { + for (int64_t nn = 0; nn < n; ++nn) { + buf_qdq_ptr[kk * n + nn] = + static_cast(v_src[kk * v_stride_n + nn]); + } + } + ::executorch::cpublas::gemm( + ::executorch::cpublas::TransposeType::NoTranspose, + ::executorch::cpublas::TransposeType::NoTranspose, + n, + m, + k, + static_cast(1), + buf_qdq_ptr, + n, + static_cast(qk_data), + qk_stride_m, + beta, + o_data, + o_stride_m); + return; + } auto do_gemm = [&](auto rt_tag) { using rt = decltype(rt_tag); ::executorch::cpublas::gemm( @@ -1207,9 +1236,17 @@ void cpu_flash_attention( // For reduced-precision activations the attention weights are cast to // the activation dtype so that Softmax(q @ k.T) @ v runs as a // reduced-in/float-out gemm matching the value matrix. + // Casting the weights down only pays when the bf16 attn@V kernel is + // the one that runs. Above kMinQBlockForWidenedAV it is faster to keep + // them in accum_t, widen V and let BLAS multiply -- also one rounding + // step fewer. Below it the widening stops amortizing. + constexpr int64_t kMinQBlockForWidenedAV = 64; + const bool widen_v = is_reduced_type && !is_quantized_sdpa && + std::is_same::value && + qBlockSize >= kMinQBlockForWidenedAV; const void* qk_gemm_data = qk_data; if constexpr (is_reduced_type) { - if (!is_quantized_sdpa) { + if (!is_quantized_sdpa && !widen_v) { vec::convert( qk_data, qk_reduced_data, qBlockSize * kvBlockSize); qk_gemm_data = qk_reduced_data; @@ -1227,7 +1264,8 @@ void cpu_flash_attention( dst_data, headSize, n == 0 ? static_cast(0) : static_cast(1), - buf_qdq_ptr); + buf_qdq_ptr, + widen_v); } // dst <- dst / sum[row] // reorder MHA output with strides From dcc197785f46dd2c855c18dce4fe325f24c38d2e Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 10 Aug 2026 16:20:43 -0400 Subject: [PATCH 09/12] [blas] Native AVX512-BF16 (vdpbf16ps) bf16 dot on x86 bf16_dot_with_fp32_arith had a native bf16-dot path only for ARM (vbfdotq_f32); x86 fell back to convert-bf16->fp32 + fp32 FMA. Add an x86 AVX512-BF16 path using _mm512_dpbf16_ps (native bf16 pairwise dot, fp32 accumulate), behind __attribute__((target(...avx512bf16))) so it compiles without TU-wide -mavx512bf16, dispatched at runtime via cpuinfo_has_x86_avx512bf16(). Identical fp32-accumulate numerics; falls back to the existing path on non-avx512bf16 hw. Used by custom SDPA's bf16 q@k.T / attn@v dots. gemma-3-1b bf16 8da4w: custom_sdpa 372->56ms, Method::execute 669->351ms (prefill @16c 670->357ms, now faster than the fp32 build's 401ms); argmax unchanged (107). Unchanged from the original commit except that the cpuinfo_initialize() fix it also carried now lands separately, before this, since it fixes the aarch64 dispatch too. --- kernels/optimized/blas/BlasKernel.cpp | 75 +++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/kernels/optimized/blas/BlasKernel.cpp b/kernels/optimized/blas/BlasKernel.cpp index 153493e32b1..c3e52c00417 100644 --- a/kernels/optimized/blas/BlasKernel.cpp +++ b/kernels/optimized/blas/BlasKernel.cpp @@ -22,6 +22,9 @@ #ifdef __aarch64__ #include #include +#elif defined(__x86_64__) +#include +#include #endif namespace vec = at::vec; @@ -116,6 +119,17 @@ float reduce(vec::VectorizedN& x) { #endif // defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && // defined(__clang__) && __clang_major__ > 15 +// GCC 10 / Clang 9 are the first versions with both the target("avx512bf16") +// function attribute and _mm512_dpbf16_ps. +#if defined(__x86_64__) && defined(__clang__) && __clang_major__ >= 9 +#define COMPILER_SUPPORTS_X86_BF16_TARGET 1 +#elif defined(__x86_64__) && !defined(__clang__) && defined(__GNUC__) && \ + __GNUC__ >= 10 +#define COMPILER_SUPPORTS_X86_BF16_TARGET 1 +#else +#define COMPILER_SUPPORTS_X86_BF16_TARGET 0 +#endif + #if COMPILER_SUPPORTS_BF16_TARGET #define TARGET_ARM_BF16_ATTRIBUTE __attribute__((target("arch=armv8.2-a+bf16"))) @@ -326,6 +340,62 @@ dot_with_fp32_arith_no_bfdot(const T* vec1, const T* vec2, int64_t len) { } #undef DOT_WITH_FP32_ARITH_TAIL_AFTER_MAIN_LOOP_BODY +#if COMPILER_SUPPORTS_X86_BF16_TARGET +// Native x86 bf16 dot using AVX512-BF16's vdpbf16ps (_mm512_dpbf16_ps), +// which computes bf16 x bf16 -> fp32 accumulate in a single instruction. +__attribute__((target("avx512f,avx512bw,avx512vl,avx512bf16"))) float +dot_with_fp32_arith_x86bfdot( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t len) { + // Each __m512bh holds 32 bf16; _mm512_dpbf16_ps accumulates bf16 x bf16 + // products into a __m512 of 16 fp32 lanes. + constexpr int kBF16PerRegister = 32; + constexpr int kAccumulators = 4; + constexpr int kBF16PerIteration = kBF16PerRegister * kAccumulators; + + __m512 acc[kAccumulators]; + for (int i = 0; i < kAccumulators; ++i) { + acc[i] = _mm512_setzero_ps(); + } + + int64_t j = 0; + // Main loop: 4 independent accumulators for ILP, 128 bf16 per iteration. + const int64_t len_main = len - (len % kBF16PerIteration); + for (; j < len_main; j += kBF16PerIteration) { + for (int i = 0; i < kAccumulators; ++i) { + const int64_t off = j + i * kBF16PerRegister; + const __m512bh a = + (__m512bh)_mm512_loadu_si512((const void*)(vec1 + off)); + const __m512bh b = + (__m512bh)_mm512_loadu_si512((const void*)(vec2 + off)); + acc[i] = _mm512_dpbf16_ps(acc[i], a, b); + } + } + + // 32-wide cleanup loop for the remaining full registers. + const int64_t len_vec = len - (len % kBF16PerRegister); + for (; j < len_vec; j += kBF16PerRegister) { + const __m512bh a = (__m512bh)_mm512_loadu_si512((const void*)(vec1 + j)); + const __m512bh b = (__m512bh)_mm512_loadu_si512((const void*)(vec2 + j)); + acc[0] = _mm512_dpbf16_ps(acc[0], a, b); + } + + float reduced_sum = 0; + for (int i = 0; i < kAccumulators; ++i) { + reduced_sum += _mm512_reduce_add_ps(acc[i]); + } + + // Scalar fp32 tail for the remainder, matching no_bfdot numerics. + for (; j < len; ++j) { + float x1 = vec1[j]; + float x2 = vec2[j]; + reduced_sum += x1 * x2; + } + return reduced_sum; +} +#endif // COMPILER_SUPPORTS_X86_BF16_TARGET + } // namespace float bf16_dot_with_fp32_arith( @@ -338,6 +408,11 @@ float bf16_dot_with_fp32_arith( return dot_with_fp32_arith_bfdot(vec1, vec2, len); } else #endif // COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_X86_BF16_TARGET + if (cpuinfo_initialize() && cpuinfo_has_x86_avx512bf16()) { + return dot_with_fp32_arith_x86bfdot(vec1, vec2, len); + } else +#endif // COMPILER_SUPPORTS_X86_BF16_TARGET { return dot_with_fp32_arith_no_bfdot(vec1, vec2, len); } From a0f3e3ca723c21a396d9991769bb91c53db6f5b6 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 11 Aug 2026 09:31:46 -0700 Subject: [PATCH 10/12] [blas] x86 counterpart of the four-column bf16 dot dot4_with_fp32_arith_bfdot gave gemm_transa_ a four-output block on aarch64. Add the AVX512-BF16 equivalent so x86 gets the same structure. The win is larger here than the aarch64 argument alone suggests. The x86 single-dot path already carries four accumulators for instruction-level parallelism, but custom SDPA calls it with len == headSize (64 for a typical LLM), which is below the 128-element main loop: it fills one accumulator and then reduces all four, so three of every four cross-lane reductions run over zeros. Giving each accumulator a distinct output makes the same four reductions do four times the work. Dispatch mirrors the single-dot path, cpuinfo_has_x86_avx512bf16 behind COMPILER_SUPPORTS_X86_BF16_TARGET, and the scalar tail keeps the numerics of the no_bfdot fallback. --- kernels/optimized/blas/BlasKernel.cpp | 51 +++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/kernels/optimized/blas/BlasKernel.cpp b/kernels/optimized/blas/BlasKernel.cpp index c3e52c00417..43fed96971a 100644 --- a/kernels/optimized/blas/BlasKernel.cpp +++ b/kernels/optimized/blas/BlasKernel.cpp @@ -457,6 +457,51 @@ TARGET_ARM_BF16_ATTRIBUTE static void dot4_with_fp32_arith_bfdot( } #endif // COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_X86_BF16_TARGET +// x86 counterpart of dot4_with_fp32_arith_bfdot. The single-dot path above +// carries four accumulators for ILP, but at len == headSize it never enters +// its 128-wide main loop: it fills one accumulator and then reduces all four, +// so three of every four reductions are over zeros. Here each accumulator +// holds a distinct output instead, so the same four reductions do four times +// the work. +__attribute__((target("avx512f,avx512bw,avx512vl,avx512bf16"))) static void +dot4_with_fp32_arith_x86bfdot( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t stride2, + int64_t len, + float* out) { + constexpr int64_t kBF16PerRegister = 32; + __m512 acc[4]; + for (int i = 0; i < 4; ++i) { + acc[i] = _mm512_setzero_ps(); + } + + int64_t j = 0; + const int64_t len_vec = len - (len % kBF16PerRegister); + for (; j < len_vec; j += kBF16PerRegister) { + const __m512bh a = (__m512bh)_mm512_loadu_si512((const void*)(vec1 + j)); + for (int i = 0; i < 4; ++i) { + const __m512bh b = + (__m512bh)_mm512_loadu_si512((const void*)(vec2 + i * stride2 + j)); + acc[i] = _mm512_dpbf16_ps(acc[i], a, b); + } + } + + for (int i = 0; i < 4; ++i) { + out[i] = _mm512_reduce_add_ps(acc[i]); + } + + // Scalar fp32 tail, matching the numerics of the no_bfdot path. + for (; j < len; ++j) { + const float x1 = vec1[j]; + for (int i = 0; i < 4; ++i) { + out[i] += x1 * static_cast(vec2[i * stride2 + j]); + } + } +} +#endif // COMPILER_SUPPORTS_X86_BF16_TARGET + void bf16_dot4_with_fp32_arith( const at::BFloat16* vec1, const at::BFloat16* vec2, @@ -469,6 +514,12 @@ void bf16_dot4_with_fp32_arith( return; } #endif // COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_X86_BF16_TARGET + if (cpuinfo_initialize() && cpuinfo_has_x86_avx512bf16()) { + dot4_with_fp32_arith_x86bfdot(vec1, vec2, stride2, len, out); + return; + } +#endif // COMPILER_SUPPORTS_X86_BF16_TARGET for (int64_t j = 0; j < 4; ++j) { out[j] = bf16_dot_with_fp32_arith(vec1, vec2 + j * stride2, len); } From defd0bd449a342d829208e36434e1bf6ecbdc93c Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 11 Aug 2026 12:14:46 -0700 Subject: [PATCH 11/12] [blas] Vectorize bf16 decode GEMV on x86 and ARM Custom SDPA decode reaches both attention GEMMs with `n == 1`, where the existing general loops leave bf16 slower than fp32. Add register-blocked attn@V kernels that keep fp32 output accumulators live while streaming bf16 inputs: AVX512 conversion plus FMA on x86, and NEON `shll` plus `fmla` on AArch64. For Q@K on x86, tile four key rows through the native AVX512-BF16 dot path. Preserve portable fallbacks, `alpha` and `beta` behavior, and vector-width tail handling. Keep the new architecture-specific implementations and dispatch together in one ARM/x86/portable conditional block. On AMD Genoa, bf16 decode rises from 19.58 to 24.77 tok/s at s2048 and from 9.30 to 17.14 tok/s at s8192, beating fp32 by 8.5% and 45.1%. Output digests are unchanged. Validated with the production `portable_lib` build, the decode GEMV correctness harness, AArch64 cross-compilation and qemu execution, clang-format, and lintrunner. Authored with Codex. --- kernels/optimized/blas/BlasKernel.cpp | 514 ++++++++++++++++++++---- kernels/optimized/blas/BlasKernel.h | 43 +- kernels/optimized/test/libblas_test.cpp | 64 +++ 3 files changed, 526 insertions(+), 95 deletions(-) diff --git a/kernels/optimized/blas/BlasKernel.cpp b/kernels/optimized/blas/BlasKernel.cpp index 43fed96971a..574d82caa7d 100644 --- a/kernels/optimized/blas/BlasKernel.cpp +++ b/kernels/optimized/blas/BlasKernel.cpp @@ -19,6 +19,8 @@ #include #include +#include + #ifdef __aarch64__ #include #include @@ -340,16 +342,180 @@ dot_with_fp32_arith_no_bfdot(const T* vec1, const T* vec2, int64_t len) { } #undef DOT_WITH_FP32_ARITH_TAIL_AFTER_MAIN_LOOP_BODY -#if COMPILER_SUPPORTS_X86_BF16_TARGET +} // namespace + +#if COMPILER_SUPPORTS_BF16_TARGET +// Four dots against a shared vec1: the cross-lane reduction is paid once per +// four results rather than per result, which matters when callers reduce over +// headSize while producing a whole tile of outputs. +TARGET_ARM_BF16_ATTRIBUTE static void dot4_with_fp32_arith_bfdot( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t stride2, + int64_t len, + float* out) { + constexpr int64_t kElementsPerIteration = 8; + float32x4_t acc[4] = { + vdupq_n_f32(0.0f), + vdupq_n_f32(0.0f), + vdupq_n_f32(0.0f), + vdupq_n_f32(0.0f)}; + int64_t idx = 0; + for (; idx + kElementsPerIteration <= len; idx += kElementsPerIteration) { + // See NOTE[Intrinsics in bfdot variant] above. + const auto v1 = vld1q_bf16(reinterpret_cast(&vec1[idx])); + for (int64_t j = 0; j < 4; ++j) { + const auto v2 = vld1q_bf16( + reinterpret_cast(&vec2[j * stride2 + idx])); + acc[j] = vbfdotq_f32(acc[j], v1, v2); + } + } + const float32x4_t sums = + vpaddq_f32(vpaddq_f32(acc[0], acc[1]), vpaddq_f32(acc[2], acc[3])); + vst1q_f32(out, sums); + for (; idx < len; ++idx) { + for (int64_t j = 0; j < 4; ++j) { + out[j] += static_cast(vec1[idx]) * + static_cast(vec2[j * stride2 + idx]); + } + } +} +#endif // COMPILER_SUPPORTS_BF16_TARGET + +#if defined(__aarch64__) +template +C10_ALWAYS_INLINE void gemv_notrans_block_neon( + int64_t output_offset, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + float32x4_t acc[kRegisterPairs * 2]; + for (int64_t r = 0; r < kRegisterPairs * 2; ++r) { + if (beta == 0.0f) { + acc[r] = vdupq_n_f32(0.0f); + } else { + acc[r] = vld1q_f32(c + output_offset + r * 4); + if (beta != 1.0f) { + acc[r] = vmulq_n_f32(acc[r], beta); + } + } + } + + for (int64_t l = 0; l < k; ++l) { + const float b_val = static_cast(b[l]) * alpha; + const auto* a_col = + reinterpret_cast(a + l * lda + output_offset); + for (int64_t r = 0; r < kRegisterPairs; ++r) { + const uint16x8_t a_bf16 = vld1q_u16(a_col + r * 8); + const float32x4_t a_low = + vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(a_bf16), 16)); + const float32x4_t a_high = + vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(a_bf16), 16)); + acc[r * 2] = vfmaq_n_f32(acc[r * 2], a_low, b_val); + acc[r * 2 + 1] = vfmaq_n_f32(acc[r * 2 + 1], a_high, b_val); + } + } + + for (int64_t r = 0; r < kRegisterPairs * 2; ++r) { + vst1q_f32(c + output_offset + r * 4, acc[r]); + } +} + +static void gemv_notrans_neon( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + int64_t i = 0; + for (; i + 64 <= m; i += 64) { + gemv_notrans_block_neon<8>(i, k, alpha, a, lda, b, beta, c); + } + for (; i + 32 <= m; i += 32) { + gemv_notrans_block_neon<4>(i, k, alpha, a, lda, b, beta, c); + } + for (; i + 16 <= m; i += 16) { + gemv_notrans_block_neon<2>(i, k, alpha, a, lda, b, beta, c); + } + for (; i + 8 <= m; i += 8) { + gemv_notrans_block_neon<1>(i, k, alpha, a, lda, b, beta, c); + } + for (; i < m; ++i) { + float acc = beta == 0.0f ? 0.0f : beta * c[i]; + for (int64_t l = 0; l < k; ++l) { + acc += static_cast(a[l * lda + i]) * + (static_cast(b[l]) * alpha); + } + c[i] = acc; + } +} + +static float +platform_bf16_dot(const BFloat16* vec1, const BFloat16* vec2, int64_t len) { +#if COMPILER_SUPPORTS_BF16_TARGET + if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { + return dot_with_fp32_arith_bfdot(vec1, vec2, len); + } +#endif // COMPILER_SUPPORTS_BF16_TARGET + return dot_with_fp32_arith_no_bfdot(vec1, vec2, len); +} + +static bool try_platform_bf16_dot4( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t stride2, + int64_t len, + float* out) { +#if COMPILER_SUPPORTS_BF16_TARGET + if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { + dot4_with_fp32_arith_bfdot(vec1, vec2, stride2, len, out); + return true; + } +#endif // COMPILER_SUPPORTS_BF16_TARGET + return false; +} + +static bool try_platform_bf16_gemv_notrans( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + gemv_notrans_neon(m, k, alpha, a, lda, b, beta, c); + return true; +} + +static bool try_platform_bf16_gemv_transa( + int64_t, + int64_t, + float, + const BFloat16*, + int64_t, + const BFloat16*, + float, + float*) { + return false; +} + +#elif COMPILER_SUPPORTS_X86_BF16_TARGET + // Native x86 bf16 dot using AVX512-BF16's vdpbf16ps (_mm512_dpbf16_ps), // which computes bf16 x bf16 -> fp32 accumulate in a single instruction. -__attribute__((target("avx512f,avx512bw,avx512vl,avx512bf16"))) float +__attribute__((target("avx512f,avx512bw,avx512vl,avx512bf16"))) static float dot_with_fp32_arith_x86bfdot( const BFloat16* vec1, const BFloat16* vec2, int64_t len) { - // Each __m512bh holds 32 bf16; _mm512_dpbf16_ps accumulates bf16 x bf16 - // products into a __m512 of 16 fp32 lanes. constexpr int kBF16PerRegister = 32; constexpr int kAccumulators = 4; constexpr int kBF16PerIteration = kBF16PerRegister * kAccumulators; @@ -360,7 +526,6 @@ dot_with_fp32_arith_x86bfdot( } int64_t j = 0; - // Main loop: 4 independent accumulators for ILP, 128 bf16 per iteration. const int64_t len_main = len - (len % kBF16PerIteration); for (; j < len_main; j += kBF16PerIteration) { for (int i = 0; i < kAccumulators; ++i) { @@ -373,7 +538,6 @@ dot_with_fp32_arith_x86bfdot( } } - // 32-wide cleanup loop for the remaining full registers. const int64_t len_vec = len - (len % kBF16PerRegister); for (; j < len_vec; j += kBF16PerRegister) { const __m512bh a = (__m512bh)_mm512_loadu_si512((const void*)(vec1 + j)); @@ -385,79 +549,14 @@ dot_with_fp32_arith_x86bfdot( for (int i = 0; i < kAccumulators; ++i) { reduced_sum += _mm512_reduce_add_ps(acc[i]); } - - // Scalar fp32 tail for the remainder, matching no_bfdot numerics. for (; j < len; ++j) { - float x1 = vec1[j]; - float x2 = vec2[j]; + const float x1 = vec1[j]; + const float x2 = vec2[j]; reduced_sum += x1 * x2; } return reduced_sum; } -#endif // COMPILER_SUPPORTS_X86_BF16_TARGET - -} // namespace - -float bf16_dot_with_fp32_arith( - const at::BFloat16* vec1, - const at::BFloat16* vec2, - int64_t len) { -#if COMPILER_SUPPORTS_BF16_TARGET - // cpuinfo_has_* reads a zeroed struct until cpuinfo_initialize() runs. - if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { - return dot_with_fp32_arith_bfdot(vec1, vec2, len); - } else -#endif // COMPILER_SUPPORTS_BF16_TARGET -#if COMPILER_SUPPORTS_X86_BF16_TARGET - if (cpuinfo_initialize() && cpuinfo_has_x86_avx512bf16()) { - return dot_with_fp32_arith_x86bfdot(vec1, vec2, len); - } else -#endif // COMPILER_SUPPORTS_X86_BF16_TARGET - { - return dot_with_fp32_arith_no_bfdot(vec1, vec2, len); - } -} -#if COMPILER_SUPPORTS_BF16_TARGET -// Four dots against a shared vec1: the cross-lane reduction is paid once per -// four results rather than per result, which matters when callers reduce over -// headSize while producing a whole tile of outputs. -TARGET_ARM_BF16_ATTRIBUTE static void dot4_with_fp32_arith_bfdot( - const BFloat16* vec1, - const BFloat16* vec2, - int64_t stride2, - int64_t len, - float* out) { - constexpr int64_t kElementsPerIteration = 8; - float32x4_t acc[4] = { - vdupq_n_f32(0.0f), - vdupq_n_f32(0.0f), - vdupq_n_f32(0.0f), - vdupq_n_f32(0.0f)}; - int64_t idx = 0; - for (; idx + kElementsPerIteration <= len; idx += kElementsPerIteration) { - // See NOTE[Intrinsics in bfdot variant] above. - const auto v1 = - vld1q_bf16(reinterpret_cast(&vec1[idx])); - for (int64_t j = 0; j < 4; ++j) { - const auto v2 = vld1q_bf16( - reinterpret_cast(&vec2[j * stride2 + idx])); - acc[j] = vbfdotq_f32(acc[j], v1, v2); - } - } - const float32x4_t sums = vpaddq_f32( - vpaddq_f32(acc[0], acc[1]), vpaddq_f32(acc[2], acc[3])); - vst1q_f32(out, sums); - for (; idx < len; ++idx) { - for (int64_t j = 0; j < 4; ++j) { - out[j] += static_cast(vec1[idx]) * - static_cast(vec2[j * stride2 + idx]); - } - } -} -#endif // COMPILER_SUPPORTS_BF16_TARGET - -#if COMPILER_SUPPORTS_X86_BF16_TARGET // x86 counterpart of dot4_with_fp32_arith_bfdot. The single-dot path above // carries four accumulators for ILP, but at len == headSize it never enters // its 128-wide main loop: it fills one accumulator and then reduces all four, @@ -500,29 +599,272 @@ dot4_with_fp32_arith_x86bfdot( } } } -#endif // COMPILER_SUPPORTS_X86_BF16_TARGET -void bf16_dot4_with_fp32_arith( - const at::BFloat16* vec1, - const at::BFloat16* vec2, +template +__attribute__(( + target("avx512f,avx512bw,avx512vl,avx512bf16,fma"))) C10_ALWAYS_INLINE void +gemv_notrans_block_x86bf16( + int64_t output_offset, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + __m512 acc[kRegisters]; + for (int64_t r = 0; r < kRegisters; ++r) { + if (beta == 0.0f) { + acc[r] = _mm512_setzero_ps(); + } else { + acc[r] = _mm512_loadu_ps(c + output_offset + r * 16); + if (beta != 1.0f) { + acc[r] = _mm512_mul_ps(acc[r], _mm512_set1_ps(beta)); + } + } + } + + for (int64_t l = 0; l < k; ++l) { + const __m512 b_vec = _mm512_set1_ps(static_cast(b[l]) * alpha); + const BFloat16* a_col = a + l * lda + output_offset; + for (int64_t r = 0; r < kRegisters; ++r) { + const __m256bh a_vec = (__m256bh)_mm256_loadu_si256( + reinterpret_cast(a_col + r * 16)); + acc[r] = _mm512_fmadd_ps(_mm512_cvtpbh_ps(a_vec), b_vec, acc[r]); + } + } + + for (int64_t r = 0; r < kRegisters; ++r) { + _mm512_storeu_ps(c + output_offset + r * 16, acc[r]); + } +} + +__attribute__((target("avx512f,avx512bw,avx512vl,avx512bf16,fma"))) static void +gemv_notrans_x86bf16( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + int64_t i = 0; + for (; i + 128 <= m; i += 128) { + gemv_notrans_block_x86bf16<8>(i, k, alpha, a, lda, b, beta, c); + } + for (; i + 64 <= m; i += 64) { + gemv_notrans_block_x86bf16<4>(i, k, alpha, a, lda, b, beta, c); + } + for (; i + 16 <= m; i += 16) { + gemv_notrans_block_x86bf16<1>(i, k, alpha, a, lda, b, beta, c); + } + for (; i < m; ++i) { + float acc = beta == 0.0f ? 0.0f : beta * c[i]; + for (int64_t l = 0; l < k; ++l) { + acc += static_cast(a[l * lda + i]) * + (static_cast(b[l]) * alpha); + } + c[i] = acc; + } +} + +__attribute__((target("avx512f,avx512bw,avx512vl,avx512bf16"))) static void +gemv_transa_x86bfdot( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + int64_t i = 0; + for (; i + 4 <= m; i += 4) { + float dots[4]; + dot4_with_fp32_arith_x86bfdot(b, a + i * lda, lda, k, dots); + for (int64_t d = 0; d < 4; ++d) { + c[i + d] = + beta == 0.0f ? alpha * dots[d] : beta * c[i + d] + alpha * dots[d]; + } + } + for (; i < m; ++i) { + const float dot = dot_with_fp32_arith_x86bfdot(a + i * lda, b, k); + c[i] = beta == 0.0f ? alpha * dot : beta * c[i] + alpha * dot; + } +} + +static bool use_x86_bf16() { + return cpuinfo_initialize() && cpuinfo_has_x86_avx512bf16(); +} + +static float +platform_bf16_dot(const BFloat16* vec1, const BFloat16* vec2, int64_t len) { + return use_x86_bf16() ? dot_with_fp32_arith_x86bfdot(vec1, vec2, len) + : dot_with_fp32_arith_no_bfdot(vec1, vec2, len); +} + +static bool try_platform_bf16_dot4( + const BFloat16* vec1, + const BFloat16* vec2, int64_t stride2, int64_t len, float* out) { -#if COMPILER_SUPPORTS_BF16_TARGET - if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { - dot4_with_fp32_arith_bfdot(vec1, vec2, stride2, len, out); - return; + if (!use_x86_bf16()) { + return false; } -#endif // COMPILER_SUPPORTS_BF16_TARGET -#if COMPILER_SUPPORTS_X86_BF16_TARGET - if (cpuinfo_initialize() && cpuinfo_has_x86_avx512bf16()) { - dot4_with_fp32_arith_x86bfdot(vec1, vec2, stride2, len, out); + dot4_with_fp32_arith_x86bfdot(vec1, vec2, stride2, len, out); + return true; +} + +static bool try_platform_bf16_gemv_notrans( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + if (!use_x86_bf16()) { + return false; + } + gemv_notrans_x86bf16(m, k, alpha, a, lda, b, beta, c); + return true; +} + +static bool try_platform_bf16_gemv_transa( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + if (!use_x86_bf16()) { + return false; + } + gemv_transa_x86bfdot(m, k, alpha, a, lda, b, beta, c); + return true; +} + +#else + +static float +platform_bf16_dot(const BFloat16* vec1, const BFloat16* vec2, int64_t len) { + return dot_with_fp32_arith_no_bfdot(vec1, vec2, len); +} + +static bool try_platform_bf16_dot4( + const BFloat16*, + const BFloat16*, + int64_t, + int64_t, + float*) { + return false; +} + +static bool try_platform_bf16_gemv_notrans( + int64_t, + int64_t, + float, + const BFloat16*, + int64_t, + const BFloat16*, + float, + float*) { + return false; +} + +static bool try_platform_bf16_gemv_transa( + int64_t, + int64_t, + float, + const BFloat16*, + int64_t, + const BFloat16*, + float, + float*) { + return false; +} + +#endif // defined(__aarch64__) + +float bf16_dot_with_fp32_arith( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t len) { + return platform_bf16_dot(vec1, vec2, len); +} + +void bf16_dot4_with_fp32_arith( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t stride2, + int64_t len, + float* out) { + if (try_platform_bf16_dot4(vec1, vec2, stride2, len, out)) { return; } -#endif // COMPILER_SUPPORTS_X86_BF16_TARGET for (int64_t j = 0; j < 4; ++j) { out[j] = bf16_dot_with_fp32_arith(vec1, vec2 + j * stride2, len); } } +void bf16_gemv_notrans_with_fp32_arith( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + if (try_platform_bf16_gemv_notrans(m, k, alpha, a, lda, b, beta, c)) { + return; + } + if (beta == 0.0f) { + std::fill(c, c + m, 0.0f); + } else if (beta != 1.0f) { + for (int64_t i = 0; i < m; ++i) { + c[i] *= beta; + } + } + for (int64_t l = 0; l < k; ++l) { + const BFloat16* a_col = a + l * lda; + const float b_val = static_cast(b[l]) * alpha; + for (int64_t i = 0; i < m; ++i) { + c[i] += static_cast(a_col[i]) * b_val; + } + } +} + +void bf16_gemv_transa_with_fp32_arith( + int64_t m, + int64_t k, + float alpha, + const BFloat16* a, + int64_t lda, + const BFloat16* b, + float beta, + float* c) { + if (try_platform_bf16_gemv_transa(m, k, alpha, a, lda, b, beta, c)) { + return; + } + int64_t i = 0; + for (; i + 4 <= m; i += 4) { + float dots[4]; + bf16_dot4_with_fp32_arith(b, a + i * lda, lda, k, dots); + for (int64_t d = 0; d < 4; ++d) { + c[i + d] = + beta == 0.0f ? alpha * dots[d] : beta * c[i + d] + alpha * dots[d]; + } + } + for (; i < m; ++i) { + const float dot = bf16_dot_with_fp32_arith(a + i * lda, b, k); + c[i] = beta == 0.0f ? alpha * dot : beta * c[i] + alpha * dot; + } +} + } // namespace executorch::cpublas::internal diff --git a/kernels/optimized/blas/BlasKernel.h b/kernels/optimized/blas/BlasKernel.h index 850bc72049e..557a1767bc6 100644 --- a/kernels/optimized/blas/BlasKernel.h +++ b/kernels/optimized/blas/BlasKernel.h @@ -146,6 +146,24 @@ void bf16_dot4_with_fp32_arith( int64_t stride2, int64_t len, float* out); +void bf16_gemv_notrans_with_fp32_arith( + int64_t m, + int64_t k, + float alpha, + const torch::executor::BFloat16* a, + int64_t lda, + const torch::executor::BFloat16* b, + float beta, + float* c); +void bf16_gemv_transa_with_fp32_arith( + int64_t m, + int64_t k, + float alpha, + const torch::executor::BFloat16* a, + int64_t lda, + const torch::executor::BFloat16* b, + float beta, + float* c); } // namespace internal // Used by custom SDPA's attn@V. Serial on purpose: SDPA already parallelizes @@ -163,16 +181,17 @@ gemm_notrans_( const torch::executor::BFloat16 *b, int64_t ldb, float beta, float *c, int64_t ldc) { + if (n == 1) { + internal::bf16_gemv_notrans_with_fp32_arith( + m, k, alpha, a, lda, b, beta, c); + return; + } + // bf16_dot_with_fp32_arith needs a contiguous k-vector, but a is strided by - // lda in k, so the dot path must gather each row first: k strided loads to - // buy n dots. At n == 1 that is one gather per multiply-add and cannot pay - // on any architecture. a's m-dimension is contiguous, so accumulate along it - // instead, which is the traversal the fp32 specialization already uses. - // - // On aarch64 the gather-free form keeps winning well past n == 1. On x86 - // vdpbf16ps makes the dot strong enough that only n == 1 avoids it. The - // crossover is empirical, and only visible with runtime-valued dimensions: - // constant-folded ones let the gather-free loop unroll and hide it. + // lda in k, so the dot path must gather each row first. On aarch64 the + // gather-free form wins for small n; x86 uses the dot path from n == 2. + // The crossover is empirical and only visible with runtime-valued + // dimensions: constant-folded ones let the gather-free loop unroll. #if defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) constexpr int64_t kMinColsForGather = 32; #else @@ -315,6 +334,12 @@ inline void gemm_transa_( const torch::executor::BFloat16 *b, int64_t ldb, float beta, float *c, int64_t ldc) { + if (n == 1) { + internal::bf16_gemv_transa_with_fp32_arith( + m, k, alpha, a, lda, b, beta, c); + return; + } + // Four columns at a time: k is headSize here, short enough that each dot's // cross-lane reduction is a large share of its cost, and this tile produces // m*n of them. diff --git a/kernels/optimized/test/libblas_test.cpp b/kernels/optimized/test/libblas_test.cpp index 358db95a672..8ee4b95014a 100644 --- a/kernels/optimized/test/libblas_test.cpp +++ b/kernels/optimized/test/libblas_test.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #define TEST_FORALL_SUPPORTED_CTYPES(_, N) \ @@ -232,6 +233,69 @@ TEST(BlasTest, BF16FloatGemmMatchesScalarAccumulation) { } } +TEST(BlasTest, BF16FloatGemmDecodeShapesMatchScalarAccumulation) { + using executorch::aten::BFloat16; + using executorch::cpublas::TransposeType; + + struct Shape { + bool transa; + int64_t m; + int64_t k; + }; + constexpr Shape kShapes[] = { + {false, 64, 511}, + {false, 128, 512}, + {false, 130, 513}, + {true, 512, 64}, + {true, 515, 128}, + {true, 513, 130}, + }; + + for (const Shape shape : kShapes) { + constexpr int64_t kN = 1; + const int64_t lda = (shape.transa ? shape.k : shape.m) + 3; + const int64_t ldb = shape.k; + const int64_t ldc = shape.m; + const auto a = + make_values(lda * (shape.transa ? shape.m : shape.k), 8); + const auto b = make_values(shape.k, 9); + + for (const auto [alpha, beta] : + {std::pair{1.0f, 0.0f}, std::pair{-0.5f, 0.25f}}) { + auto c = make_values(shape.m, 10); + auto expected = c; + + // clang-format off + reference_gemm( + shape.transa, + shape.m, kN, shape.k, + alpha, + a.data(), lda, + b.data(), ldb, + beta, + expected, ldc); + + executorch::cpublas::gemm( + shape.transa ? TransposeType::Transpose : TransposeType::NoTranspose, + TransposeType::NoTranspose, + shape.m, kN, shape.k, + alpha, + a.data(), lda, + b.data(), ldb, + beta, + c.data(), ldc); + // clang-format on + + const std::string context = "transa=" + std::to_string(shape.transa) + + " m=" + std::to_string(shape.m) + " k=" + std::to_string(shape.k) + + " alpha=" + std::to_string(alpha) + " beta=" + std::to_string(beta); + for (int64_t i = 0; i < shape.m; ++i) { + expect_near_relative(c[i], expected[i], context.c_str()); + } + } + } +} + // beta == 0 must overwrite c rather than read it, so uninitialized garbage in // the output cannot poison the result. TEST(BlasTest, BF16FloatGemmBetaZeroIgnoresOutput) { From 779b3b574ad482777ceeedb1216f026a8aedcb3d Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 11 Aug 2026 12:29:57 -0700 Subject: [PATCH 12/12] [blas] Clarify ARM bf16 target macro Rename the ARM-only compiler capability macro to distinguish it from the corresponding x86 bf16 target check. Validated with clang-format and lintrunner. Authored with Codex. --- kernels/optimized/blas/BlasKernel.cpp | 30 +++++++++++++-------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/kernels/optimized/blas/BlasKernel.cpp b/kernels/optimized/blas/BlasKernel.cpp index 574d82caa7d..6b514251b7d 100644 --- a/kernels/optimized/blas/BlasKernel.cpp +++ b/kernels/optimized/blas/BlasKernel.cpp @@ -109,15 +109,15 @@ float reduce(vec::VectorizedN& x) { #if defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && \ defined(__clang__) && __clang_major__ > 15 // https://godbolt.org/z/z8P4Yncra -#define COMPILER_SUPPORTS_BF16_TARGET 1 +#define COMPILER_SUPPORTS_ARM_BF16_TARGET 1 #elif defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && \ !defined(__clang__) && defined(__GNUC__) && __GNUC__ >= 10 // https://gcc.gnu.org/gcc-10/changes.html // https://godbolt.org/z/cdGG7vn8o -#define COMPILER_SUPPORTS_BF16_TARGET 1 +#define COMPILER_SUPPORTS_ARM_BF16_TARGET 1 #else // defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && // defined(__clang__) && __clang_major__ > 15 -#define COMPILER_SUPPORTS_BF16_TARGET 0 +#define COMPILER_SUPPORTS_ARM_BF16_TARGET 0 #endif // defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE) && // defined(__clang__) && __clang_major__ > 15 @@ -132,7 +132,7 @@ float reduce(vec::VectorizedN& x) { #define COMPILER_SUPPORTS_X86_BF16_TARGET 0 #endif -#if COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_ARM_BF16_TARGET #define TARGET_ARM_BF16_ATTRIBUTE __attribute__((target("arch=armv8.2-a+bf16"))) TARGET_ARM_BF16_ATTRIBUTE C10_ALWAYS_INLINE void @@ -171,7 +171,7 @@ dot_with_fp32_arith_vectorized_tail_inner_loop_bfdot( #else #define TARGET_ARM_BF16_ATTRIBUTE -#endif // COMPILER_SUPPORTS_BF16_TARGET +#endif // COMPILER_SUPPORTS_ARM_BF16_TARGET namespace { @@ -249,7 +249,7 @@ C10_ALWAYS_INLINE auto dot_with_fp32_arith_main_loop_no_bfdot( return reduce(sum); } -#if COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_ARM_BF16_TARGET template struct ForcedUnrollTargetBFloat16 { template @@ -287,7 +287,7 @@ dot_with_fp32_arith_main_loop_bfdot( } return reduce(sum); } -#endif // COMPILER_SUPPORTS_BF16_TARGET +#endif // COMPILER_SUPPORTS_ARM_BF16_TARGET static_assert( (vec::Vectorized::size() & @@ -324,7 +324,7 @@ static_assert( } \ return reduced_sum -#if COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_ARM_BF16_TARGET TARGET_ARM_BF16_ATTRIBUTE float dot_with_fp32_arith_bfdot( const BFloat16* vec1, const BFloat16* vec2, @@ -332,7 +332,7 @@ TARGET_ARM_BF16_ATTRIBUTE float dot_with_fp32_arith_bfdot( auto reduced_sum = dot_with_fp32_arith_main_loop_bfdot(vec1, vec2, len); DOT_WITH_FP32_ARITH_TAIL_AFTER_MAIN_LOOP_BODY(_bfdot); } -#endif // COMPILER_SUPPORTS_BF16_TARGET +#endif // COMPILER_SUPPORTS_ARM_BF16_TARGET template C10_ALWAYS_INLINE float @@ -344,7 +344,7 @@ dot_with_fp32_arith_no_bfdot(const T* vec1, const T* vec2, int64_t len) { } // namespace -#if COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_ARM_BF16_TARGET // Four dots against a shared vec1: the cross-lane reduction is paid once per // four results rather than per result, which matters when callers reduce over // headSize while producing a whole tile of outputs. @@ -380,7 +380,7 @@ TARGET_ARM_BF16_ATTRIBUTE static void dot4_with_fp32_arith_bfdot( } } } -#endif // COMPILER_SUPPORTS_BF16_TARGET +#endif // COMPILER_SUPPORTS_ARM_BF16_TARGET #if defined(__aarch64__) template @@ -459,11 +459,11 @@ static void gemv_notrans_neon( static float platform_bf16_dot(const BFloat16* vec1, const BFloat16* vec2, int64_t len) { -#if COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_ARM_BF16_TARGET if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { return dot_with_fp32_arith_bfdot(vec1, vec2, len); } -#endif // COMPILER_SUPPORTS_BF16_TARGET +#endif // COMPILER_SUPPORTS_ARM_BF16_TARGET return dot_with_fp32_arith_no_bfdot(vec1, vec2, len); } @@ -473,12 +473,12 @@ static bool try_platform_bf16_dot4( int64_t stride2, int64_t len, float* out) { -#if COMPILER_SUPPORTS_BF16_TARGET +#if COMPILER_SUPPORTS_ARM_BF16_TARGET if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { dot4_with_fp32_arith_bfdot(vec1, vec2, stride2, len, out); return true; } -#endif // COMPILER_SUPPORTS_BF16_TARGET +#endif // COMPILER_SUPPORTS_ARM_BF16_TARGET return false; }