diff --git a/compression/q4_0-inl.h b/compression/q4_0-inl.h index 97526251..a5a68c4f 100644 --- a/compression/q4_0-inl.h +++ b/compression/q4_0-inl.h @@ -97,25 +97,57 @@ class Q4_0Codec { template > static HWY_INLINE void DequantizeBlock(D d, const uint8_t* HWY_RESTRICT block_ptr, Raw* HWY_RESTRICT raw) { - const hn::Repartition df; - const hn::Rebind di32; - const hn::Rebind di16; - const hn::Rebind di8; - const hn::Rebind du8; + const hn::Full128 du8_128; + const hn::Full128 di8_128; + const hn::Full128 di16_128; + const hn::Full128 di32_128; + const hn::Full128 df_128; using T = ScaleT; T scale; hwy::CopyBytes(block_ptr, &scale, sizeof(T)); const float scale_f = hwy::F32FromBF16(scale); - const auto vd = hn::Set(df, scale_f); + const auto vd_128 = hn::Set(df_128, scale_f); const uint8_t* qs_ptr = block_ptr + sizeof(T); - const size_t N = hn::Lanes(df); - const size_t num_vectors = 32 / N; - - for (size_t v_idx = 0; v_idx < num_vectors; ++v_idx) { - const auto out = DequantizeLanes(df, du8, di8, qs_ptr, 0, v_idx * N, N, vd); - StoreRaw(df, out, raw + v_idx * N); + const auto raw_16 = hn::LoadU(du8_128, qs_ptr); + const auto mask_0f = hn::Set(du8_128, 0x0F); + const auto offset_8 = hn::Set(di8_128, 8); + + const auto low_i8 = + hn::Sub(hn::BitCast(di8_128, hn::And(raw_16, mask_0f)), offset_8); + const auto high_i8 = + hn::Sub(hn::BitCast(di8_128, hn::ShiftRight<4>(raw_16)), offset_8); + + const auto low_i16_0 = hn::PromoteLowerTo(di16_128, low_i8); + const auto low_i16_1 = hn::PromoteUpperTo(di16_128, low_i8); + const auto high_i16_0 = hn::PromoteLowerTo(di16_128, high_i8); + const auto high_i16_1 = hn::PromoteUpperTo(di16_128, high_i8); + + const auto f0 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, low_i16_0)), vd_128); + const auto f1 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, low_i16_0)), vd_128); + const auto f2 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, low_i16_1)), vd_128); + const auto f3 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, low_i16_1)), vd_128); + const auto f4 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, high_i16_0)), vd_128); + const auto f5 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, high_i16_0)), vd_128); + const auto f6 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteLowerTo(di32_128, high_i16_1)), vd_128); + const auto f7 = hn::Mul(hn::ConvertTo(df_128, hn::PromoteUpperTo(di32_128, high_i16_1)), vd_128); + + if constexpr (hwy::IsSame()) { + hn::StoreU(f0, df_128, raw + 0); + hn::StoreU(f1, df_128, raw + 4); + hn::StoreU(f2, df_128, raw + 8); + hn::StoreU(f3, df_128, raw + 12); + hn::StoreU(f4, df_128, raw + 16); + hn::StoreU(f5, df_128, raw + 20); + hn::StoreU(f6, df_128, raw + 24); + hn::StoreU(f7, df_128, raw + 28); + } else { + const hn::Full128 dbf_128; + hn::StoreU(hn::OrderedDemote2To(dbf_128, f0, f1), dbf_128, raw + 0); + hn::StoreU(hn::OrderedDemote2To(dbf_128, f2, f3), dbf_128, raw + 8); + hn::StoreU(hn::OrderedDemote2To(dbf_128, f4, f5), dbf_128, raw + 16); + hn::StoreU(hn::OrderedDemote2To(dbf_128, f6, f7), dbf_128, raw + 24); } } @@ -259,18 +291,24 @@ class Q4_0Codec { HWY_DASSERT(current_packed_ofs % kBlockSize == 0); const size_t num_full_blocks = num_to_decompress / kBlockSize; - for (size_t b = 0; b < num_full_blocks; ++b) { - const uint8_t* block_ptr = - &packed.ptr->byte + BlockByteOffset(current_packed_ofs); + const uint8_t* block_ptr = + &packed.ptr->byte + BlockByteOffset(current_packed_ofs); + size_t b = 0; + for (; b + 1 < num_full_blocks; b += 2) { + DequantizeBlock(d, block_ptr, current_raw); + DequantizeBlock(d, block_ptr + 18, current_raw + 32); + block_ptr += 36; + current_raw += 64; + } + if (b < num_full_blocks) { DequantizeBlock(d, block_ptr, current_raw); - current_packed_ofs += kBlockSize; - current_raw += kBlockSize; + block_ptr += 18; + current_raw += 32; } + current_packed_ofs += num_full_blocks * kBlockSize; const size_t remaining = num_to_decompress % kBlockSize; if (remaining != 0) { - const uint8_t* block_ptr = - &packed.ptr->byte + BlockByteOffset(current_packed_ofs); HWY_ALIGN Raw temp[kBlockSize]; DequantizeBlock(d, block_ptr, temp); memcpy(current_raw, temp, remaining * sizeof(Raw)); diff --git a/compression/test_util-inl.h b/compression/test_util-inl.h index bb2fadb0..9e11d385 100644 --- a/compression/test_util-inl.h +++ b/compression/test_util-inl.h @@ -119,9 +119,15 @@ MatStorageT GenerateMat(const Extents2D& extents, MatPadding padding, f = -f; // Also generate some negative values. row[c] = f; } - Compress(raw.Row(r), raw.Cols(), ws.tls[thread], - MakeSpan(compressed.Row(r), extents.cols), - /*packed_ofs=*/0); + if constexpr (IsQ4_0Stream()) { + Compress(raw.Row(r), raw.Cols(), ws.tls[thread], + compressed.Span(), + /*packed_ofs=*/r * extents.cols); + } else { + Compress(raw.Row(r), raw.Cols(), ws.tls[thread], + MakeSpan(compressed.Row(r), extents.cols), + /*packed_ofs=*/0); + } }); compressed.SetScale(0.6f); // Arbitrary value, different from 1. @@ -148,9 +154,15 @@ MatStorageT GenerateTransposedMat(const Extents2D extents, f = -f; // Also generate some negative values. row[c] = f; } - Compress(raw.Row(r), raw.Cols(), ws.tls[thread], - MakeSpan(compressed.Row(r), extents.cols), - /*packed_ofs=*/0); + if constexpr (IsQ4_0Stream()) { + Compress(raw.Row(r), raw.Cols(), ws.tls[thread], + compressed.Span(), + /*packed_ofs=*/r * extents.cols); + } else { + Compress(raw.Row(r), raw.Cols(), ws.tls[thread], + MakeSpan(compressed.Row(r), extents.cols), + /*packed_ofs=*/0); + } }); // Arbitrary value, different from 1, must match `GenerateMat`. @@ -205,15 +217,24 @@ void AssertClose(const MatPtrT& A, const MatPtrT& B, MatStorageT c_slow_batch("c_slow_batch", Extents2D(A.Rows(), B_rows), allocator, MatPadding::kOdd); for (size_t m = 0; m < A.Rows(); ++m) { - DecompressAndZeroPad(df, MakeSpan(A.Row(m), cols), 0, a_batch.Row(m), cols); + if constexpr (IsQ4_0Stream()) { + DecompressAndZeroPad(df, A.PaddedSpan(), m * cols, a_batch.Row(m), cols); + } else { + DecompressAndZeroPad(df, MakeSpan(A.Row(m), cols), 0, a_batch.Row(m), cols); + } DecompressAndZeroPad(df, MakeSpan(C.Row(m), B_rows), 0, c_batch.Row(m), B_rows); DecompressAndZeroPad(df, MakeSpan(C_slow.Row(m), B_rows), 0, c_slow_batch.Row(m), B_rows); } for (size_t n = 0; n < B_rows; ++n) { - DecompressAndZeroPad(df, MakeSpan(B.Row(n), cols), 0, b_trans_batch.Row(n), - cols); + if constexpr (IsQ4_0Stream()) { + DecompressAndZeroPad(df, B.PaddedSpan(), n * cols, b_trans_batch.Row(n), + cols); + } else { + DecompressAndZeroPad(df, MakeSpan(B.Row(n), cols), 0, b_trans_batch.Row(n), + cols); + } } // MatMul rounds inputs to BF16, so error is proportional to the max input @@ -231,6 +252,9 @@ void AssertClose(const MatPtrT& A, const MatPtrT& B, if (IsF32() || IsF32()) { tolerance += 2 * max_abs * eps_bf16; } + if constexpr (IsQ4_0Stream()) { + tolerance += 0.02 * norm; + } if (tolerance > 500.0) { HWY_WARN("high tolerance %f norm %f maxabs %f\n", tolerance, norm, max_abs); diff --git a/evals/gemma_batch_bench.cc b/evals/gemma_batch_bench.cc index ea5e9793..6b886718 100644 --- a/evals/gemma_batch_bench.cc +++ b/evals/gemma_batch_bench.cc @@ -36,15 +36,11 @@ GemmaEnv* s_env = nullptr; class GemmaBatchBench : public ::testing::Test { protected: - std::vector BatchGemmaReply( + QueryResultAndMetrics BatchGemmaReplyWithMetrics( const std::vector& inputs) { s_env->MutableConfig().temperature = 0.0f; // deterministic s_env->MutableConfig().verbosity = 2; - std::vector replies; - for (const QueryResult& result : s_env->BatchQueryModel(inputs)) { - replies.push_back(result.response); - } - return replies; + return s_env->BatchQueryModelWithMetrics(inputs); } }; @@ -128,16 +124,99 @@ std::vector GenerateInputs() { TEST_F(GemmaBatchBench, RandomQuestionsBatched) { s_env->SetMaxGeneratedTokens(12); const std::vector inputs = GenerateInputs(); - // Run multiple times so that auto-tuning is closer to complete. - for (size_t rep = 0; rep < 4; ++rep) { - std::vector responses = BatchGemmaReply(inputs); + constexpr size_t kNumReps = 7; + + std::vector prefill_speeds; + std::vector generate_speeds; + size_t total_prefill_tokens = 0; + double total_prefill_duration = 0.0; + size_t total_generate_tokens = 0; + double total_generate_duration = 0.0; + + size_t warm_prefill_tokens = 0; + double warm_prefill_duration = 0.0; + size_t warm_generate_tokens = 0; + double warm_generate_duration = 0.0; + + for (size_t rep = 0; rep < kNumReps; ++rep) { + QueryResultAndMetrics result = BatchGemmaReplyWithMetrics(inputs); + const std::vector& responses = result.query_results; + const TimingInfo& timing = result.timing_info; + + const double prefill_tok_sec = + timing.prefill_duration > 0.0 + ? static_cast(timing.prefill_tokens) / + timing.prefill_duration + : 0.0; + const double gen_tok_sec = + timing.generate_duration > 0.0 + ? static_cast(timing.tokens_generated) / + timing.generate_duration + : 0.0; + + prefill_speeds.push_back(prefill_tok_sec); + generate_speeds.push_back(gen_tok_sec); + + total_prefill_tokens += timing.prefill_tokens; + total_prefill_duration += timing.prefill_duration; + total_generate_tokens += timing.tokens_generated; + total_generate_duration += timing.generate_duration; + + if (rep > 0) { + warm_prefill_tokens += timing.prefill_tokens; + warm_prefill_duration += timing.prefill_duration; + warm_generate_tokens += timing.tokens_generated; + warm_generate_duration += timing.generate_duration; + } + for (size_t i = 0; i < HWY_MIN(hwy::Unpredictable1() * 3, responses.size()); ++i) { fprintf(stderr, "Rep %zu batch answer %zu '%s'\n\n", rep, i, - responses[i].c_str()); + responses[i].response.c_str()); } PROFILER_PRINT_RESULTS(); } + + const double avg_prefill = + total_prefill_duration > 0.0 + ? static_cast(total_prefill_tokens) / total_prefill_duration + : 0.0; + const double avg_generate = + total_generate_duration > 0.0 + ? static_cast(total_generate_tokens) / total_generate_duration + : 0.0; + + const double warm_avg_prefill = + warm_prefill_duration > 0.0 + ? static_cast(warm_prefill_tokens) / warm_prefill_duration + : 0.0; + const double warm_avg_generate = + warm_generate_duration > 0.0 + ? static_cast(warm_generate_tokens) / warm_generate_duration + : 0.0; + + fprintf(stderr, + "\n============================================================\n"); + fprintf(stderr, + "[ Gemma Batch Benchmark Summary (%zu Repetitions) ]\n", kNumReps); + for (size_t rep = 0; rep < kNumReps; ++rep) { + fprintf(stderr, + " Rep %zu: Prefill = %7.2f tok/s | Generate = %7.2f tok/s%s\n", + rep, prefill_speeds[rep], generate_speeds[rep], + rep == 0 ? " (warmup / autotune)" : ""); + } + fprintf(stderr, + "------------------------------------------------------------\n"); + fprintf(stderr, + "Overall Average: Prefill = %7.2f tok/s | Generate = %7.2f tok/s\n", + avg_prefill, avg_generate); + if (kNumReps > 1) { + fprintf(stderr, + "Warm Average : Prefill = %7.2f tok/s | Generate = %7.2f tok/s\n", + warm_avg_prefill, warm_avg_generate); + } + fprintf(stderr, + "============================================================\n\n"); } } // namespace diff --git a/ops/bench_matmul.cc b/ops/bench_matmul.cc index 6616cf3d..250d908a 100644 --- a/ops/bench_matmul.cc +++ b/ops/bench_matmul.cc @@ -184,15 +184,19 @@ void BenchAllMatMul() { // QKV projection BenchMatMul(batch_size, 1152, 1536, kAdd, env); BenchMatMul(batch_size, 1152, 1536, kAdd, env); + BenchMatMul(batch_size, 1152, 1536, kAdd, env); // FFN gate+up BenchMatMul(batch_size, 1152, 13824, kAdd, env); BenchMatMul(batch_size, 1152, 13824, kAdd, env); + BenchMatMul(batch_size, 1152, 13824, kAdd, env); // FFN down BenchMatMul(batch_size, 6912, 1152, kAdd, env); BenchMatMul(batch_size, 6912, 1152, kAdd, env); + BenchMatMul(batch_size, 6912, 1152, kAdd, env); // Logits / embedding BenchMatMul(batch_size, 1152, 262144, kAdd, env); BenchMatMul(batch_size, 1152, 262144, kAdd, env); + BenchMatMul(batch_size, 1152, 262144, kAdd, env); } PROFILER_PRINT_RESULTS(); diff --git a/ops/matmul-inl.h b/ops/matmul-inl.h index 449c9517..7437cd51 100644 --- a/ops/matmul-inl.h +++ b/ops/matmul-inl.h @@ -403,6 +403,15 @@ class MMKernel { const IndexRange& range_mc, const IndexRange& range_kc, const IndexRange& range_nc, const MMArgs& args, Tag out_tag, CView C_MC_NC) { + if constexpr (IsQ4_0Stream()) { + if (HWY_LIKELY(range_mc.Num() <= 4 && range_kc.begin() % 32 == 0 && + B.Stride() % 32 == 0)) { + B3A2C0_Q4_0(A, B, range_mc, range_kc, range_nc, args, out_tag, + C_MC_NC); + return; + } + } + const size_t kc = range_kc.Num(); const StridedViewBF A_view = A.View(range_mc.begin(), range_kc.begin(), kc); @@ -782,6 +791,414 @@ class MMKernel { } HWY_DASSERT(imc == mc); } + +#if !HWY_TARGET_IS_SVE && !HWY_HAVE_SCALABLE + struct UnpackedBlockQ4_0 { + HWY_ATTR UnpackedBlockQ4_0() = default; + HWY_ATTR UnpackedBlockQ4_0(hn::Vec> l, + hn::Vec> h) + : low(l), high(h) {} + hn::Vec> low; + hn::Vec> high; + }; + + static HWY_INLINE HWY_ATTR UnpackedBlockQ4_0 UnpackBlockQ4_0( + const uint8_t* HWY_RESTRICT blk, const hn::Full128 du8, + const hn::Full128 di8, + const hn::Vec> mask_0f, + const hn::Vec> offset_8) { + const auto raw = hn::LoadU(du8, blk + 2); + return {hn::Sub(hn::BitCast(di8, hn::And(raw, mask_0f)), offset_8), + hn::Sub(hn::BitCast(di8, hn::ShiftRight<4>(raw)), offset_8)}; + } +#else + static HWY_INLINE HWY_ATTR void UnpackBlockQ4_0( + const uint8_t* HWY_RESTRICT blk, const hn::Full128 du8, + const hn::Full128 di8, + const hn::Vec> mask_0f, + const hn::Vec> offset_8, + hn::Vec>& low, + hn::Vec>& high) { + const auto raw = hn::LoadU(du8, blk + 2); + low = hn::Sub(hn::BitCast(di8, hn::And(raw, mask_0f)), offset_8); + high = hn::Sub(hn::BitCast(di8, hn::ShiftRight<4>(raw)), offset_8); + } +#endif + + static HWY_INLINE HWY_ATTR float QuantizeActivationRow( + const BF16* HWY_RESTRICT ar, size_t kc, size_t num_blocks, + int8_t* HWY_RESTRICT q_a_row) { + const hn::Full128 dbf; + const hn::Full128 df; + const hn::Full128 di32; + const hn::Full128 di16; + const hn::Full128 di8; + + auto v_max = hn::Zero(df); + size_t k = 0; + for (; k + 8 <= kc; k += 8) { + const auto bf = hn::LoadU(dbf, ar + k); + const auto f0 = hn::PromoteLowerTo(df, bf); + const auto f1 = hn::PromoteUpperTo(df, bf); + v_max = hn::Max(v_max, hn::Max(hn::Abs(f0), hn::Abs(f1))); + } + float max_abs = hn::ReduceMax(df, v_max); + for (; k < kc; ++k) { + max_abs = std::max(max_abs, std::abs(hwy::F32FromBF16(ar[k]))); + } + + if (HWY_UNLIKELY(max_abs == 0.0f)) { + hwy::ZeroBytes(q_a_row, num_blocks * 32); + return 0.0f; + } + + const float sa = max_abs / 127.0f; + const auto vinv_scale = hn::Set(df, 127.0f / max_abs); + for (size_t b = 0; b < num_blocks; ++b) { + const size_t block_len = HWY_MIN(kc - b * 32, 32); + const BF16* HWY_RESTRICT ab = ar + b * 32; + int8_t* HWY_RESTRICT qb = q_a_row + b * 32; + if (HWY_LIKELY(block_len == 32)) { + const auto bf0 = hn::LoadU(dbf, ab + 0); + const auto bf1 = hn::LoadU(dbf, ab + 8); + const auto bf2 = hn::LoadU(dbf, ab + 16); + const auto bf3 = hn::LoadU(dbf, ab + 24); + + const auto f0 = hn::PromoteLowerTo(df, bf0); + const auto f1 = hn::PromoteUpperTo(df, bf0); + const auto f2 = hn::PromoteLowerTo(df, bf1); + const auto f3 = hn::PromoteUpperTo(df, bf1); + const auto f4 = hn::PromoteLowerTo(df, bf2); + const auto f5 = hn::PromoteUpperTo(df, bf2); + const auto f6 = hn::PromoteLowerTo(df, bf3); + const auto f7 = hn::PromoteUpperTo(df, bf3); + + const auto i0 = hn::NearestInt(hn::Mul(f0, vinv_scale)); + const auto i1 = hn::NearestInt(hn::Mul(f1, vinv_scale)); + const auto i2 = hn::NearestInt(hn::Mul(f2, vinv_scale)); + const auto i3 = hn::NearestInt(hn::Mul(f3, vinv_scale)); + const auto i4 = hn::NearestInt(hn::Mul(f4, vinv_scale)); + const auto i5 = hn::NearestInt(hn::Mul(f5, vinv_scale)); + const auto i6 = hn::NearestInt(hn::Mul(f6, vinv_scale)); + const auto i7 = hn::NearestInt(hn::Mul(f7, vinv_scale)); + + const auto p16_0 = hn::OrderedDemote2To(di16, i0, i1); + const auto p16_1 = hn::OrderedDemote2To(di16, i2, i3); + const auto p16_2 = hn::OrderedDemote2To(di16, i4, i5); + const auto p16_3 = hn::OrderedDemote2To(di16, i6, i7); + + hn::Store(hn::OrderedDemote2To(di8, p16_0, p16_1), di8, qb + 0); + hn::Store(hn::OrderedDemote2To(di8, p16_2, p16_3), di8, qb + 16); + } else { + const float inv_scale = 127.0f / max_abs; + HWY_ALIGN int8_t temp[32] = {}; + for (size_t i = 0; i < block_len; ++i) { + const float v0 = hwy::F32FromBF16(ab[i]) * inv_scale; + temp[i] = static_cast(std::min( + 127, std::max(-128, static_cast(std::round(v0))))); + } + memcpy(qb, temp, 32); + } + } + return sa; + } + + static HWY_INLINE HWY_ATTR hn::Vec> TransposeReduce4( + const hn::Full128 di32, const hn::Full128 df, + const hn::Vec> acc0, + const hn::Vec> acc1, + const hn::Vec> acc2, + const hn::Vec> acc3) { + const auto ab_sum = hn::Add(hn::InterleaveLower(di32, acc0, acc1), + hn::InterleaveUpper(di32, acc0, acc1)); + const auto cd_sum = hn::Add(hn::InterleaveLower(di32, acc2, acc3), + hn::InterleaveUpper(di32, acc2, acc3)); + const auto acbd_sum = hn::Add(hn::InterleaveLower(di32, ab_sum, cd_sum), + hn::InterleaveUpper(di32, ab_sum, cd_sum)); + return hn::ConvertTo(df, acbd_sum); + } + + template + static HWY_INLINE HWY_ATTR void LoopKC_Q4_0( + const StridedViewBF A_view, const MatPtrT& B, + const PackedSpan& B_span, size_t imc, size_t kc, + size_t num_blocks, size_t col0, const IndexRange& range_nc, + const float scale, const float* HWY_RESTRICT add, Tag tag, + CView C_MC_NC) { + HWY_DASSERT(num_blocks <= kMaxKC / 32); + + const hn::Full128 du8; + const hn::Full128 di8; + const hn::Full128 dbf; + const hn::Full128 di32; + const hn::Full128 df; + + float scale_a[kRowsAC]; + HWY_ALIGN int8_t q_a[kRowsAC][kMaxKC]; + for (size_t r = 0; r < kRowsAC; ++r) { + scale_a[r] = QuantizeActivationRow(A_view.Row(imc + r), kc, num_blocks, + q_a[r]); + } + + const auto mask_0f = hn::Set(du8, 0x0F); + const auto offset_8 = hn::Set(di8, 8); + + const uint8_t* HWY_RESTRICT B_base = &B_span.ptr->byte; + const size_t b_stride = B.Stride(); + + for (size_t inc = 0; inc < range_nc.Num(); inc += kNR) { + const size_t row_b = range_nc.begin() + inc; + const CView C_MC_NR = C_MC_NC.View(0, inc, kNR); + const float* HWY_RESTRICT add_row = add ? add + row_b : nullptr; + + const uint8_t* ptr_b0 = + B_base + ((row_b + 0) * b_stride + col0) / 32 * 18; + const uint8_t* ptr_b1 = + B_base + ((row_b + 1) * b_stride + col0) / 32 * 18; + const uint8_t* ptr_b2 = + B_base + ((row_b + 2) * b_stride + col0) / 32 * 18; + const uint8_t* ptr_b3 = + B_base + ((row_b + 3) * b_stride + col0) / 32 * 18; + +#if !HWY_TARGET_IS_SVE && !HWY_HAVE_SCALABLE + hn::Vec sum[kRowsAC]; + for (size_t r = 0; r < kRowsAC; ++r) { + sum[r] = hn::Zero(df); + } + + for (size_t b = 0; b < num_blocks; ++b) { + hwy::Prefetch(ptr_b0 + (b + 4) * 18); + hwy::Prefetch(ptr_b1 + (b + 4) * 18); + hwy::Prefetch(ptr_b2 + (b + 4) * 18); + hwy::Prefetch(ptr_b3 + (b + 4) * 18); + + const uint8_t* blk0 = ptr_b0 + b * 18; + const uint8_t* blk1 = ptr_b1 + b * 18; + const uint8_t* blk2 = ptr_b2 + b * 18; + const uint8_t* blk3 = ptr_b3 + b * 18; + + hwy::bfloat16_t sb0_bf, sb1_bf, sb2_bf, sb3_bf; + hwy::CopyBytes(blk0, &sb0_bf, 2); + hwy::CopyBytes(blk1, &sb1_bf, 2); + hwy::CopyBytes(blk2, &sb2_bf, 2); + hwy::CopyBytes(blk3, &sb3_bf, 2); + + const auto v_bf16 = hn::Dup128VecFromValues( + dbf, sb0_bf, sb2_bf, sb1_bf, sb3_bf, hwy::bfloat16_t(), + hwy::bfloat16_t(), hwy::bfloat16_t(), hwy::bfloat16_t()); + const auto v_sb = hn::PromoteLowerTo(df, v_bf16); + + const auto b0 = UnpackBlockQ4_0(blk0, du8, di8, mask_0f, offset_8); + const auto b1 = UnpackBlockQ4_0(blk1, du8, di8, mask_0f, offset_8); + const auto b2 = UnpackBlockQ4_0(blk2, du8, di8, mask_0f, offset_8); + const auto b3 = UnpackBlockQ4_0(blk3, du8, di8, mask_0f, offset_8); + + for (size_t r = 0; r < kRowsAC; ++r) { + if (scale_a[r] == 0.0f) continue; + + const auto a_low = hn::Load(di8, q_a[r] + b * 32); + const auto a_high = hn::Load(di8, q_a[r] + b * 32 + 16); + + const auto acc0 = hn::SumOfMulQuadAccumulate( + di32, b0.high, a_high, + hn::SumOfMulQuadAccumulate(di32, b0.low, a_low, hn::Zero(di32))); + const auto acc1 = hn::SumOfMulQuadAccumulate( + di32, b1.high, a_high, + hn::SumOfMulQuadAccumulate(di32, b1.low, a_low, hn::Zero(di32))); + const auto acc2 = hn::SumOfMulQuadAccumulate( + di32, b2.high, a_high, + hn::SumOfMulQuadAccumulate(di32, b2.low, a_low, hn::Zero(di32))); + const auto acc3 = hn::SumOfMulQuadAccumulate( + di32, b3.high, a_high, + hn::SumOfMulQuadAccumulate(di32, b3.low, a_low, hn::Zero(di32))); + + const auto dot_f = + TransposeReduce4(di32, df, acc0, acc1, acc2, acc3); + sum[r] = hn::MulAdd(dot_f, v_sb, sum[r]); + } + } + + for (size_t r = 0; r < kRowsAC; ++r) { + const auto s_scaled = hn::Mul(sum[r], hn::Set(df, scale_a[r])); + const auto s0 = hn::GetLane(s_scaled); + const auto s2 = hn::GetLane(hn::ShiftRightLanes<1>(df, s_scaled)); + const auto s1 = hn::GetLane(hn::ShiftRightLanes<2>(df, s_scaled)); + const auto s3 = hn::GetLane(hn::ShiftRightLanes<3>(df, s_scaled)); + sum[r] = hn::Dup128VecFromValues(df, s0, s1, s2, s3); + } + + MMStoreHorizontalSumsIntoC horz; + const hn::Vec sum0 = sum[0]; + const hn::Vec sum1 = + kRowsAC > 1 ? sum[1] : hn::Zero(df); + const hn::Vec sum2 = + kRowsAC > 2 ? sum[2] : hn::Zero(df); + const hn::Vec sum3 = + kRowsAC > 3 ? sum[3] : hn::Zero(df); + + horz.Store(df, sum0, sum1, sum2, sum3, scale, add_row, imc, tag, + C_MC_NR); +#else + hn::Vec sum0 = hn::Zero(df); + hn::Vec sum1 = hn::Zero(df); + hn::Vec sum2 = hn::Zero(df); + hn::Vec sum3 = hn::Zero(df); + + for (size_t b = 0; b < num_blocks; ++b) { + hwy::Prefetch(ptr_b0 + (b + 4) * 18); + hwy::Prefetch(ptr_b1 + (b + 4) * 18); + hwy::Prefetch(ptr_b2 + (b + 4) * 18); + hwy::Prefetch(ptr_b3 + (b + 4) * 18); + + const uint8_t* blk0 = ptr_b0 + b * 18; + const uint8_t* blk1 = ptr_b1 + b * 18; + const uint8_t* blk2 = ptr_b2 + b * 18; + const uint8_t* blk3 = ptr_b3 + b * 18; + + hwy::bfloat16_t sb0_bf, sb1_bf, sb2_bf, sb3_bf; + hwy::CopyBytes(blk0, &sb0_bf, 2); + hwy::CopyBytes(blk1, &sb1_bf, 2); + hwy::CopyBytes(blk2, &sb2_bf, 2); + hwy::CopyBytes(blk3, &sb3_bf, 2); + + const auto v_bf16 = hn::Dup128VecFromValues( + dbf, sb0_bf, sb2_bf, sb1_bf, sb3_bf, hwy::bfloat16_t(), + hwy::bfloat16_t(), hwy::bfloat16_t(), hwy::bfloat16_t()); + const auto v_sb = hn::PromoteLowerTo(df, v_bf16); + + hn::Vec b0_low, b0_high; + hn::Vec b1_low, b1_high; + hn::Vec b2_low, b2_high; + hn::Vec b3_low, b3_high; + UnpackBlockQ4_0(blk0, du8, di8, mask_0f, offset_8, b0_low, b0_high); + UnpackBlockQ4_0(blk1, du8, di8, mask_0f, offset_8, b1_low, b1_high); + UnpackBlockQ4_0(blk2, du8, di8, mask_0f, offset_8, b2_low, b2_high); + UnpackBlockQ4_0(blk3, du8, di8, mask_0f, offset_8, b3_low, b3_high); + + auto accumulate_row = [&](size_t r, auto& sum_r) HWY_ATTR { + if (scale_a[r] == 0.0f) return; + + const auto a_low = hn::Load(di8, q_a[r] + b * 32); + const auto a_high = hn::Load(di8, q_a[r] + b * 32 + 16); + + const auto acc0 = hn::SumOfMulQuadAccumulate( + di32, b0_high, a_high, + hn::SumOfMulQuadAccumulate(di32, b0_low, a_low, hn::Zero(di32))); + const auto acc1 = hn::SumOfMulQuadAccumulate( + di32, b1_high, a_high, + hn::SumOfMulQuadAccumulate(di32, b1_low, a_low, hn::Zero(di32))); + const auto acc2 = hn::SumOfMulQuadAccumulate( + di32, b2_high, a_high, + hn::SumOfMulQuadAccumulate(di32, b2_low, a_low, hn::Zero(di32))); + const auto acc3 = hn::SumOfMulQuadAccumulate( + di32, b3_high, a_high, + hn::SumOfMulQuadAccumulate(di32, b3_low, a_low, hn::Zero(di32))); + + const auto dot_f = + TransposeReduce4(di32, df, acc0, acc1, acc2, acc3); + sum_r = hn::MulAdd(dot_f, v_sb, sum_r); + }; + + accumulate_row(0, sum0); + if constexpr (kRowsAC > 1) accumulate_row(1, sum1); + if constexpr (kRowsAC > 2) accumulate_row(2, sum2); + if constexpr (kRowsAC > 3) accumulate_row(3, sum3); + } + + auto finish_row = [&](size_t r, auto s) HWY_ATTR { + const auto s_scaled = hn::Mul(s, hn::Set(df, scale_a[r])); + const auto s0 = hn::GetLane(s_scaled); + const auto s2 = hn::GetLane(hn::ShiftRightLanes<1>(df, s_scaled)); + const auto s1 = hn::GetLane(hn::ShiftRightLanes<2>(df, s_scaled)); + const auto s3 = hn::GetLane(hn::ShiftRightLanes<3>(df, s_scaled)); + return hn::Dup128VecFromValues(df, s0, s1, s2, s3); + }; + + sum0 = finish_row(0, sum0); + if constexpr (kRowsAC > 1) sum1 = finish_row(1, sum1); + if constexpr (kRowsAC > 2) sum2 = finish_row(2, sum2); + if constexpr (kRowsAC > 3) sum3 = finish_row(3, sum3); + + MMStoreHorizontalSumsIntoC horz; + horz.Store(df, sum0, sum1, sum2, sum3, scale, add_row, imc, tag, + C_MC_NR); +#endif + } + } + + template + static HWY_INLINE HWY_ATTR void A2C0_Q4_0( + const StridedViewBF A_view, const MatPtrT& B, + const PackedSpan& B_span, size_t mr, + const IndexRange& range_mc, size_t kc, size_t num_blocks, size_t col0, + const IndexRange& range_nc, const float scale, + const float* HWY_RESTRICT add, Tag tag, CView C_MC_NC) { + const size_t mc = range_mc.Num(); + size_t imc = 0; + + if (HWY_UNLIKELY(mr == 1)) { + for (; imc < mc; ++imc) { + LoopKC_Q4_0<1>(A_view, B, B_span, imc, kc, num_blocks, col0, range_nc, + scale, add, tag, C_MC_NC); + } + return; + } + + if (HWY_UNLIKELY(mr == 2)) { + if (HWY_LIKELY(mc >= 2)) { + for (; imc <= mc - 2; imc += 2) { + LoopKC_Q4_0<2>(A_view, B, B_span, imc, kc, num_blocks, col0, range_nc, + scale, add, tag, C_MC_NC); + } + } + if (HWY_UNLIKELY(imc != mc)) { + LoopKC_Q4_0<1>(A_view, B, B_span, imc, kc, num_blocks, col0, range_nc, + scale, add, tag, C_MC_NC); + } + return; + } + + HWY_DASSERT(mr == 4); + if (HWY_LIKELY(mc >= 4)) { + for (; imc <= mc - 4; imc += 4) { + LoopKC_Q4_0<4>(A_view, B, B_span, imc, kc, num_blocks, col0, range_nc, + scale, add, tag, C_MC_NC); + } + } + const size_t remainder_mc = mc - imc; + HWY_DASSERT(remainder_mc < 4); + if (HWY_UNLIKELY(remainder_mc & 2)) { + LoopKC_Q4_0<2>(A_view, B, B_span, imc, kc, num_blocks, col0, range_nc, + scale, add, tag, C_MC_NC); + imc += 2; + } + if (HWY_UNLIKELY(remainder_mc & 1)) { + LoopKC_Q4_0<1>(A_view, B, B_span, imc, kc, num_blocks, col0, range_nc, + scale, add, tag, C_MC_NC); + imc += 1; + } + HWY_DASSERT(imc == mc); + } + + template + static HWY_ATTR void B3A2C0_Q4_0(const StridedViewBF A, + const MatPtrT& B, + const IndexRange& range_mc, + const IndexRange& range_kc, + const IndexRange& range_nc, + const MMArgs& args, + Tag out_tag, + CView C_MC_NC) { + const size_t kc = range_kc.Num(); + const size_t num_blocks = hwy::DivCeil(kc, 32); + const StridedViewBF A_view = A.View(range_mc.begin(), range_kc.begin(), kc); + const PackedSpan B_span = B.PaddedSpan(); + const size_t col0 = range_kc.begin(); + const float scale = args.scale_A * B.Scale(); + + A2C0_Q4_0(A_view, B, B_span, args.mr, range_mc, kc, num_blocks, col0, + range_nc, scale, args.add, out_tag, C_MC_NC); + } }; // Miscellaneous stateless helper functions. diff --git a/ops/matmul_test.cc b/ops/matmul_test.cc index 5ae9e366..bd72c44d 100644 --- a/ops/matmul_test.cc +++ b/ops/matmul_test.cc @@ -260,6 +260,8 @@ void TestAllMatMul() { TestMatMul(256, 256, 256, /*add=*/false, env, __LINE__); TestMatMul(256, 256, 256, /*add=*/true, env, __LINE__); + TestMatMul(256, 256, 256, /*add=*/false, env, __LINE__); + TestMatMul(256, 256, 256, /*add=*/true, env, __LINE__); // Non-vector-multiple K. TestMatMul(128, 258, 128, /*add=*/true, env, __LINE__); @@ -272,30 +274,40 @@ void TestAllMatMul() { TestMatMul(33, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(31, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(29, 128, 32, /*add=*/true, env, __LINE__); + TestMatMul(31, 128, 32, /*add=*/false, env, __LINE__); + TestMatMul(29, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(4, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(4, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(4, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(4, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(4, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(4, 128, 32, /*add=*/false, env, __LINE__); + TestMatMul(4, 128, 32, /*add=*/true, env, __LINE__); + TestMatMul(4, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(3, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(3, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(3, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(3, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(3, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(3, 128, 32, /*add=*/true, env, __LINE__); + TestMatMul(3, 128, 32, /*add=*/false, env, __LINE__); + TestMatMul(3, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(2, 128, 64, /*add=*/true, env, __LINE__); TestMatMul(2, 128, 64, /*add=*/false, env, __LINE__); TestMatMul(2, 128, 64, /*add=*/true, env, __LINE__); TestMatMul(2, 128, 64, /*add=*/false, env, __LINE__); TestMatMul(2, 128, 64, /*add=*/true, env, __LINE__); TestMatMul(2, 128, 64, /*add=*/false, env, __LINE__); + TestMatMul(2, 128, 64, /*add=*/true, env, __LINE__); + TestMatMul(2, 128, 64, /*add=*/false, env, __LINE__); TestMatMul(1, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(1, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(1, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(1, 128, 32, /*add=*/true, env, __LINE__); TestMatMul(1, 128, 32, /*add=*/false, env, __LINE__); TestMatMul(1, 128, 32, /*add=*/true, env, __LINE__); + TestMatMul(1, 128, 32, /*add=*/false, env, __LINE__); + TestMatMul(1, 128, 32, /*add=*/true, env, __LINE__); pools.MaybeStopSpinning(threading_args.spin); }