From d92b46d61bc220c1edfd8e42628f72263945458f Mon Sep 17 00:00:00 2001 From: Phil Culliton Date: Tue, 25 Aug 2026 12:06:34 -0700 Subject: [PATCH] Internal change. PiperOrigin-RevId: 970706010 --- compression/q4_0-inl.h | 60 ++++--- compression/test_util-inl.h | 3 + compression/types.h | 3 +- evals/gemma_batch_bench.cc | 102 +++++++++-- ops/bench_matmul.cc | 6 +- ops/matmul-inl.h | 334 ++++++++++++++++++++++++++++++++++++ ops/matmul_test.cc | 12 ++ 7 files changed, 485 insertions(+), 35 deletions(-) diff --git a/compression/q4_0-inl.h b/compression/q4_0-inl.h index 97526251..42e2cff5 100644 --- a/compression/q4_0-inl.h +++ b/compression/q4_0-inl.h @@ -44,12 +44,12 @@ namespace HWY_NAMESPACE { namespace hn = hwy::HWY_NAMESPACE; class Q4_0Codec { - using ScaleT = hwy::bfloat16_t; - static constexpr size_t kBlockSize = 32; + using ScaleT = Q4_0Stream::ScaleT; + static constexpr size_t kBlockSize = Q4_0Stream::kBlockSize; + static constexpr size_t kBlockBytes = Q4_0Stream::kBlockBytes; static constexpr size_t BlockByteOffset(size_t packed_ofs) { - const size_t kBytesPerBlock = sizeof(ScaleT) + kBlockSize / 2; - return (packed_ofs / kBlockSize) * kBytesPerBlock; + return (packed_ofs / kBlockSize) * kBlockBytes; } template > @@ -97,25 +97,33 @@ 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::CappedTag df; + const hn::Rebind du8; const hn::Rebind di32; - const hn::Rebind di16; - const hn::Rebind di8; - const hn::Rebind du8; + const size_t N = hn::Lanes(df); 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 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 v_scale = hn::Set(df, scale_f); + const auto v_offset = hn::Set(df, scale_f * 8.0f); + + const hn::Full128 du8_16; + const auto raw_16 = hn::LoadU(du8_16, block_ptr + sizeof(T)); + const auto mask_0f = hn::Set(du8_16, 0x0F); + + HWY_ALIGN uint8_t u8_buf[2][16]; + hn::Store(hn::And(raw_16, mask_0f), du8_16, u8_buf[0]); + hn::Store(hn::ShiftRight<4>(raw_16), du8_16, u8_buf[1]); + + for (size_t half = 0; half < 2; ++half) { + for (size_t i = 0; i < 16; i += N) { + const auto u8 = hn::Load(du8, u8_buf[half] + i); + const auto u32 = hn::PromoteTo(di32, u8); + const auto val_f = hn::MulSub(hn::ConvertTo(df, u32), v_scale, v_offset); + StoreRaw(df, val_f, raw + half * 16 + i); + } } } @@ -259,18 +267,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); - current_packed_ofs += kBlockSize; + DequantizeBlock(d, block_ptr + kBlockBytes, current_raw + kBlockSize); + block_ptr += 2 * kBlockBytes; + current_raw += 2 * kBlockSize; + } + if (b < num_full_blocks) { + DequantizeBlock(d, block_ptr, current_raw); + block_ptr += kBlockBytes; current_raw += kBlockSize; } + 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 8e26cdc1..f6d52db9 100644 --- a/compression/test_util-inl.h +++ b/compression/test_util-inl.h @@ -254,6 +254,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/compression/types.h b/compression/types.h index 9454412f..f653d2c9 100644 --- a/compression/types.h +++ b/compression/types.h @@ -117,10 +117,11 @@ struct I8Stream { struct Q4_0Stream { static constexpr size_t kBlockSize = 32; using ScaleT = hwy::bfloat16_t; + static constexpr size_t kBlockBytes = sizeof(ScaleT) + kBlockSize / 2; static constexpr size_t PackedEnd(size_t capacity) { const size_t num_blocks = hwy::DivCeil(capacity, kBlockSize); - return num_blocks * (sizeof(ScaleT) + kBlockSize / 2); + return num_blocks * kBlockBytes; } uint8_t byte; diff --git a/evals/gemma_batch_bench.cc b/evals/gemma_batch_bench.cc index ea5e9793..681f0a70 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,102 @@ 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; + prefill_speeds.reserve(kNumReps); + std::vector generate_speeds; + generate_speeds.reserve(kNumReps); + + 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..11dc096c 100644 --- a/ops/bench_matmul.cc +++ b/ops/bench_matmul.cc @@ -96,7 +96,7 @@ void BenchMatMul(size_t M, size_t K, size_t N, bool add, MatMulEnv& env) { MatStorageT a = GenerateMat(A_extents, MatPadding::kOdd, env.ctx); MatStorageT b_trans = - GenerateTransposedMat(B_extents, MatPadding::kOdd, env.ctx); + GenerateTransposedMat(B_extents, MatPadding::kPacked, env.ctx); const float* add_row = add ? add_storage.PackedScale1() : nullptr; @@ -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..cabe5392 100644 --- a/ops/matmul-inl.h +++ b/ops/matmul-inl.h @@ -403,6 +403,20 @@ 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()) { + // Fast path: direct INT8 x INT4 quantized dot-product kernel for small M + // (e.g. token generation) with block-aligned K and row strides. + // If unaligned or M > 4 (e.g. prefill), falls through to the generic + // fallback path below (on-the-fly decompress of B to BF16 + BF16 matmul). + if (HWY_LIKELY(range_mc.Num() <= 4 && + range_kc.begin() % Q4_0Stream::kBlockSize == 0 && + B.Stride() % Q4_0Stream::kBlockSize == 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 +796,326 @@ class MMKernel { } HWY_DASSERT(imc == mc); } + + 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 + sizeof(Q4_0Stream::ScaleT)); + 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); + } + + 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::CappedTag dbf; + const hn::Repartition df; + const hn::Repartition di32; + const hn::Repartition di16; + const hn::Repartition di8; + const size_t N_bf = hn::Lanes(dbf); + + auto v_max = hn::Zero(df); + size_t k = 0; + for (; k + N_bf <= kc; k += N_bf) { + 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))); + } + if (HWY_UNLIKELY(k < kc)) { + const auto bf = hn::LoadN(dbf, ar + k, kc - 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))); + } + const float max_abs = hn::ReduceMax(df, v_max); + + static constexpr size_t kBlockSize = Q4_0Stream::kBlockSize; + if (HWY_UNLIKELY(max_abs == 0.0f)) { + hwy::ZeroBytes(q_a_row, num_blocks * kBlockSize); + 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 * kBlockSize, kBlockSize); + const BF16* HWY_RESTRICT ab = ar + b * kBlockSize; + int8_t* HWY_RESTRICT qb = q_a_row + b * kBlockSize; + if (HWY_LIKELY(block_len == kBlockSize)) { + for (size_t i = 0; i < kBlockSize; i += hn::Lanes(di8)) { + const auto bf0 = hn::LoadU(dbf, ab + i); + const auto bf1 = hn::LoadU(dbf, ab + i + N_bf); + + 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 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 p16_0 = hn::OrderedDemote2To(di16, i0, i1); + const auto p16_1 = hn::OrderedDemote2To(di16, i2, i3); + + hn::StoreU(hn::OrderedDemote2To(di8, p16_0, p16_1), di8, qb + i); + } + } else { + const float inv_scale = 127.0f / max_abs; + HWY_ALIGN int8_t temp[kBlockSize] = {}; + 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, kBlockSize); + } + } + 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) { + using VI32 = hn::Vec; + const VI32 ab_sum = hn::Add(hn::InterleaveLower(di32, acc0, acc1), + hn::InterleaveUpper(di32, acc0, acc1)); + const VI32 cd_sum = hn::Add(hn::InterleaveLower(di32, acc2, acc3), + hn::InterleaveUpper(di32, acc2, acc3)); + const VI32 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) { + static constexpr size_t kBlockSize = Q4_0Stream::kBlockSize; + static constexpr size_t kBlockBytes = Q4_0Stream::kBlockBytes; + using ScaleT = Q4_0Stream::ScaleT; + + HWY_DASSERT(num_blocks <= kMaxKC / kBlockSize); + + const hn::Full128 du8; + const hn::Full128 di8; + const hn::Full128 dbf; + const hn::Full128 di32; + const hn::Full128 df; + + using VU8 = hn::Vec; + using VI8 = hn::Vec; + using VBF = hn::Vec; + using VI32 = hn::Vec; + using VF = hn::Vec; + + 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 VU8 mask_0f = hn::Set(du8, 0x0F); + const VI8 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) / kBlockSize * kBlockBytes; + const uint8_t* ptr_b1 = + B_base + ((row_b + 1) * b_stride + col0) / kBlockSize * kBlockBytes; + const uint8_t* ptr_b2 = + B_base + ((row_b + 2) * b_stride + col0) / kBlockSize * kBlockBytes; + const uint8_t* ptr_b3 = + B_base + ((row_b + 3) * b_stride + col0) / kBlockSize * kBlockBytes; + + VF sum0 = hn::Zero(df); + VF sum1 = hn::Zero(df); + VF sum2 = hn::Zero(df); + VF sum3 = hn::Zero(df); + + for (size_t b = 0; b < num_blocks; ++b) { + hwy::Prefetch(ptr_b0 + (b + 4) * kBlockBytes); + hwy::Prefetch(ptr_b1 + (b + 4) * kBlockBytes); + hwy::Prefetch(ptr_b2 + (b + 4) * kBlockBytes); + hwy::Prefetch(ptr_b3 + (b + 4) * kBlockBytes); + + const uint8_t* blk0 = ptr_b0 + b * kBlockBytes; + const uint8_t* blk1 = ptr_b1 + b * kBlockBytes; + const uint8_t* blk2 = ptr_b2 + b * kBlockBytes; + const uint8_t* blk3 = ptr_b3 + b * kBlockBytes; + + hwy::bfloat16_t sb0_bf, sb1_bf, sb2_bf, sb3_bf; + hwy::CopyBytes(blk0, &sb0_bf, sizeof(ScaleT)); + hwy::CopyBytes(blk1, &sb1_bf, sizeof(ScaleT)); + hwy::CopyBytes(blk2, &sb2_bf, sizeof(ScaleT)); + hwy::CopyBytes(blk3, &sb3_bf, sizeof(ScaleT)); + + const VBF 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 VF v_sb = hn::PromoteLowerTo(df, v_bf16); + + VI8 b0_low, b0_high; + VI8 b1_low, b1_high; + VI8 b2_low, b2_high; + VI8 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, VF& sum_r) HWY_ATTR { + if (scale_a[r] == 0.0f) return; + + const VI8 a_low = hn::Load(di8, q_a[r] + b * kBlockSize); + const VI8 a_high = + hn::Load(di8, q_a[r] + b * kBlockSize + kBlockSize / 2); + + const VI32 zero = hn::Zero(di32); + const VI32 low0 = + hn::SumOfMulQuadAccumulate(di32, b0_low, a_low, zero); + const VI32 low1 = + hn::SumOfMulQuadAccumulate(di32, b1_low, a_low, zero); + const VI32 low2 = + hn::SumOfMulQuadAccumulate(di32, b2_low, a_low, zero); + const VI32 low3 = + hn::SumOfMulQuadAccumulate(di32, b3_low, a_low, zero); + + const VI32 acc0 = + hn::SumOfMulQuadAccumulate(di32, b0_high, a_high, low0); + const VI32 acc1 = + hn::SumOfMulQuadAccumulate(di32, b1_high, a_high, low1); + const VI32 acc2 = + hn::SumOfMulQuadAccumulate(di32, b2_high, a_high, low2); + const VI32 acc3 = + hn::SumOfMulQuadAccumulate(di32, b3_high, a_high, low3); + + const VF 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, VF s) HWY_ATTR { + const VF s_scaled = hn::Mul(s, hn::Set(df, scale_a[r])); + const float s0 = hn::GetLane(s_scaled); + const float s2 = hn::ExtractLane(s_scaled, 1); + const float s1 = hn::ExtractLane(s_scaled, 2); + const float s3 = hn::ExtractLane(s_scaled, 3); + 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); + } + } + + 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 790948b9..88988333 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__); #if GEMMA_ENABLE_NUQ using NUQ = NuqStream; @@ -288,30 +290,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); }