diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index 9c479569505..72fd8a60008 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( @@ -289,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); @@ -339,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( @@ -848,6 +916,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 +982,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 +1080,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 @@ -1139,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; @@ -1159,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 diff --git a/kernels/optimized/blas/BlasKernel.cpp b/kernels/optimized/blas/BlasKernel.cpp index 970105c859e..6b514251b7d 100644 --- a/kernels/optimized/blas/BlasKernel.cpp +++ b/kernels/optimized/blas/BlasKernel.cpp @@ -19,9 +19,14 @@ #include #include +#include + #ifdef __aarch64__ #include #include +#elif defined(__x86_64__) +#include +#include #endif namespace vec = at::vec; @@ -104,19 +109,30 @@ 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 -#if COMPILER_SUPPORTS_BF16_TARGET +// 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_ARM_BF16_TARGET #define TARGET_ARM_BF16_ATTRIBUTE __attribute__((target("arch=armv8.2-a+bf16"))) TARGET_ARM_BF16_ATTRIBUTE C10_ALWAYS_INLINE void @@ -155,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 { @@ -233,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 @@ -271,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() & @@ -308,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, @@ -316,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 @@ -328,17 +344,526 @@ dot_with_fp32_arith_no_bfdot(const T* vec1, const T* vec2, int64_t len) { } // namespace +#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. +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_ARM_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_ARM_BF16_TARGET + if (cpuinfo_initialize() && cpuinfo_has_arm_bf16()) { + return dot_with_fp32_arith_bfdot(vec1, vec2, len); + } +#endif // COMPILER_SUPPORTS_ARM_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_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_ARM_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"))) static float +dot_with_fp32_arith_x86bfdot( + const BFloat16* vec1, + const BFloat16* vec2, + int64_t len) { + 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; + 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); + } + } + + 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]); + } + for (; j < len; ++j) { + const float x1 = vec1[j]; + const float x2 = vec2[j]; + reduced_sum += x1 * x2; + } + return reduced_sum; +} + +// 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]); + } + } +} + +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 (!use_x86_bf16()) { + return false; + } + 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 at::BFloat16* vec1, - const at::BFloat16* vec2, + const BFloat16* vec1, + const BFloat16* vec2, int64_t len) { -#if COMPILER_SUPPORTS_BF16_TARGET - if (cpuinfo_has_arm_bf16()) { - return dot_with_fp32_arith_bfdot(vec1, vec2, len); - } else -#endif // COMPILER_SUPPORTS_BF16_TARGET - { - return dot_with_fp32_arith_no_bfdot(vec1, vec2, 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; + } + 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; } } diff --git a/kernels/optimized/blas/BlasKernel.h b/kernels/optimized/blas/BlasKernel.h index 9d4a9f87ec0..557a1767bc6 100644 --- a/kernels/optimized/blas/BlasKernel.h +++ b/kernels/optimized/blas/BlasKernel.h @@ -15,6 +15,7 @@ #include #include +#include namespace executorch { namespace cpublas { @@ -138,8 +139,121 @@ 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); +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 +// 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) { + 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. 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 + constexpr int64_t kMinColsForGather = 2; +#endif + 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) { + 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 +323,45 @@ 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) { + 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. + const auto *a_ = a; + 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; + } +} // clang-format on template 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( diff --git a/kernels/optimized/test/libblas_test.cpp b/kernels/optimized/test/libblas_test.cpp index 3c944214417..8ee4b95014a 100644 --- a/kernels/optimized/test/libblas_test.cpp +++ b/kernels/optimized/test/libblas_test.cpp @@ -11,6 +11,11 @@ #include #include +#include +#include +#include +#include +#include #include #define TEST_FORALL_SUPPORTED_CTYPES(_, N) \ @@ -45,6 +50,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 +147,185 @@ 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()); + } + } + } + } + } + } +} + +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) { + 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; + } + } +}