third_party/highway/hwy/contrib/matvec/matvec_test.cc - aom - Git at Google

 // Copyright 2023 Google LLC
 // SPDX-License-Identifier: Apache-2.0
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 #include "third_party/highway/hwy/base.h"

 // Reduce targets to avoid timeout under emulation.
 #ifndef HWY_DISABLED_TARGETS
 #define HWY_DISABLED_TARGETS (HWY_SVE2_128 | HWY_SVE2 | HWY_SVE_256 | HWY_NEON)
 #endif

 #include <stddef.h>
 #include <stdint.h>

 #include <cmath>  // std::abs

 #include "third_party/highway/hwy/aligned_allocator.h"

 // clang-format off
 #undef HWY_TARGET_INCLUDE
 #define HWY_TARGET_INCLUDE "hwy/contrib/matvec/matvec_test.cc"  // NOLINT
 #include "third_party/highway/hwy/foreach_target.h"  // IWYU pragma: keep
 // Must come after foreach_target.h
 #include "third_party/highway/hwy/contrib/algo/transform-inl.h"
 #include "third_party/highway/hwy/contrib/matvec/matvec-inl.h"
 #include "third_party/highway/hwy/highway.h"
 #include "third_party/highway/hwy/contrib/thread_pool/thread_pool.h"
 #include "third_party/highway/hwy/contrib/thread_pool/topology.h"
 #include "third_party/highway/hwy/tests/test_util-inl.h"
 // clang-format on

 HWY_BEFORE_NAMESPACE();
 namespace hwy {
 namespace HWY_NAMESPACE {
 namespace {

 template <typename MatT, typename T>
 HWY_NOINLINE void SimpleMatVecAdd(const MatT* HWY_RESTRICT mat,
                                   const T* HWY_RESTRICT vec,
                                   const T* HWY_RESTRICT add, size_t rows,
                                   size_t cols, T* HWY_RESTRICT out,
                                   ThreadPool& pool) {
   if (add) {
     pool.Run(0, rows, [=](uint64_t r, size_t /*thread*/) {
       double dot = 0.0;
       for (size_t c = 0; c < cols; c++) {
         // For reasons unknown, fp16 += does not compile on clang (Arm).
         dot += ConvertScalarTo<double>(mat[r * cols + c]) *
                ConvertScalarTo<double>(vec[c]);
       }
       out[r] = ConvertScalarTo<T>(dot + ConvertScalarTo<double>(add[r]));
     });
   } else {
     pool.Run(0, rows, [=](uint64_t r, size_t /*thread*/) {
       double dot = 0.0;
       for (size_t c = 0; c < cols; c++) {
         // For reasons unknown, fp16 += does not compile on clang (Arm).
         dot += ConvertScalarTo<double>(mat[r * cols + c]) *
                ConvertScalarTo<double>(vec[c]);
       }
       out[r] = ConvertScalarTo<T>(dot);
     });
   }
 }

 HWY_MAYBE_UNUSED HWY_NOINLINE void SimpleMatVecAdd(
     const hwy::bfloat16_t* HWY_RESTRICT mat, const float* HWY_RESTRICT vec,
     const float* add, size_t rows, size_t cols, float* HWY_RESTRICT out,
     ThreadPool& pool) {
   if (add) {
     pool.Run(0, rows, [=](uint64_t r, size_t /*thread*/) {
       float dot = 0.0f;
       for (size_t c = 0; c < cols; c++) {
         dot += F32FromBF16(mat[r * cols + c]) * vec[c];
       }
       out[r] = dot + add[r];
     });
   } else {
     pool.Run(0, rows, [=](uint64_t r, size_t /*thread*/) {
       float dot = 0.0f;
       for (size_t c = 0; c < cols; c++) {
         dot += F32FromBF16(mat[r * cols + c]) * vec[c];
       }
       out[r] = dot;
     });
   }
 }

 HWY_MAYBE_UNUSED HWY_NOINLINE void SimpleMatVecAdd(
     const hwy::bfloat16_t* HWY_RESTRICT mat,
     const hwy::bfloat16_t* HWY_RESTRICT vec,
     const hwy::bfloat16_t* HWY_RESTRICT add, size_t rows, size_t cols,
     float* HWY_RESTRICT out, ThreadPool& pool) {
   if (add) {
     pool.Run(0, rows, [=](uint64_t r, size_t /*thread*/) {
       float dot = 0.0f;
       for (size_t c = 0; c < cols; c++) {
         dot += F32FromBF16(mat[r * cols + c]) * F32FromBF16(vec[c]);
       }
       out[r] = dot + F32FromBF16(add[r]);
     });
   } else {
     pool.Run(0, rows, [=](uint64_t r, size_t /*thread*/) {
       float dot = 0.0f;
       for (size_t c = 0; c < cols; c++) {
         dot += F32FromBF16(mat[r * cols + c]) * F32FromBF16(vec[c]);
       }
       out[r] = dot;
     });
   }
 }

 // Workaround for incorrect codegen on Arm, which results in values of `av`
 // >= 1E10. Can also be prevented by calling `Print(du, indices)`.
 #if HWY_ARCH_ARM && HWY_COMPILER_CLANG
 #define GENERATE_INLINE HWY_NOINLINE
 #else
 #define GENERATE_INLINE HWY_INLINE
 #endif

 struct GenerateMod {
   template <class D, HWY_IF_NOT_BF16_D(D), HWY_IF_LANES_GT_D(D, 1)>
   GENERATE_INLINE Vec<D> operator()(D d,
                                     Vec<RebindToUnsigned<D>> indices) const {
     const RebindToUnsigned<D> du;
     return Reverse2(d, ConvertTo(d, And(indices, Set(du, 0xF))));
   }

   template <class D, HWY_IF_NOT_BF16_D(D), HWY_IF_LANES_LE_D(D, 1)>
   GENERATE_INLINE Vec<D> operator()(D d,
                                     Vec<RebindToUnsigned<D>> indices) const {
     const RebindToUnsigned<D> du;
     return ConvertTo(d, And(indices, Set(du, 0xF)));
   }

   // Requires >= 4 bf16 lanes for float32 Reverse2.
   template <class D, HWY_IF_BF16_D(D), HWY_IF_LANES_GT_D(D, 2)>
   GENERATE_INLINE Vec<D> operator()(D d,
                                     Vec<RebindToUnsigned<D>> indices) const {
     const RebindToUnsigned<D> du;
     const RebindToSigned<D> di;
     const RepartitionToWide<decltype(di)> dw;
     const RebindToFloat<decltype(dw)> df;
     indices = And(indices, Set(du, 0xF));
     const Vec<decltype(df)> i0 = ConvertTo(df, PromoteLowerTo(dw, indices));
     const Vec<decltype(df)> i1 = ConvertTo(df, PromoteUpperTo(dw, indices));
     return OrderedDemote2To(d, Reverse2(df, i0), Reverse2(df, i1));
   }

   // For one or two lanes, we don't have OrderedDemote2To nor Reverse2.
   template <class D, HWY_IF_BF16_D(D), HWY_IF_LANES_LE_D(D, 2)>
   GENERATE_INLINE Vec<D> operator()(D d,
                                     Vec<RebindToUnsigned<D>> indices) const {
     const Rebind<float, D> df;
     return DemoteTo(d, Set(df, static_cast<float>(GetLane(indices))));
   }
 };

 // MatT is usually the same as T, but can also be bfloat16_t when T = float.
 template <typename MatT, typename VecT>
 class TestMatVecAdd {
   template <size_t kRows, size_t kCols, class D, typename T = TFromD<D>>
   HWY_NOINLINE void Test(D d, ThreadPool& pool) {
 // This target lacks too many ops required in our implementation, use
 // HWY_EMU128 instead.
 #if HWY_TARGET != HWY_SCALAR
     const Repartition<MatT, D> dm;
     const Repartition<VecT, D> dv;
     const size_t misalign = 3 * Lanes(d) / 5;
     // Fill matrix and vector with small integer values
     const size_t area = kRows * kCols;
     AlignedFreeUniquePtr<MatT[]> storage_m =
         AllocateAligned<MatT>(misalign + area);
     AlignedFreeUniquePtr<VecT[]> storage_v =
         AllocateAligned<VecT>(misalign + kCols);
     AlignedFreeUniquePtr<VecT[]> storage_a =
         AllocateAligned<VecT>(misalign + kRows);
     HWY_ASSERT(storage_m && storage_v && storage_a);
     MatT* pm = storage_m.get() + misalign;
     VecT* pv = storage_v.get() + misalign;
     VecT* av = storage_a.get() + misalign;
     Generate(dm, pm, area, GenerateMod());
     Generate(dv, pv, kCols, GenerateMod());
     Generate(dv, av, kRows, GenerateMod());

     AlignedFreeUniquePtr<T[]> expected_without_add = AllocateAligned<T>(kRows);
     HWY_ASSERT(expected_without_add);
     SimpleMatVecAdd(pm, pv, static_cast<VecT*>(nullptr), kRows, kCols,
                     expected_without_add.get(), pool);

     AlignedFreeUniquePtr<T[]> actual_without_add = AllocateAligned<T>(kRows);
     HWY_ASSERT(actual_without_add);
     MatVec<kRows, kCols>(pm, pv, actual_without_add.get(), pool);

     const auto assert_close = [&](const AlignedFreeUniquePtr<T[]>& expected,
                                   const AlignedFreeUniquePtr<T[]>& actual,
                                   bool with_add) {
       for (size_t i = 0; i < kRows; ++i) {
         const double exp = ConvertScalarTo<double>(expected[i]);
         const double act = ConvertScalarTo<double>(actual[i]);
         const double epsilon =
             1.0 / (1ULL << HWY_MIN(MantissaBits<MatT>(), MantissaBits<VecT>()));
         const double tolerance = exp * 20.0 / epsilon;
         const double l1 = std::abs(exp - act);
         const double rel = exp == 0.0 ? 0.0 : l1 / exp;

         if (l1 > tolerance && rel > epsilon) {
           fprintf(stderr,
                   "%s/%s %zu x %zu, %s: mismatch at %zu: %E != %E; "
                   "tol %f l1 %f rel %E\n",
                   TypeName(MatT(), 1).c_str(), TypeName(VecT(), 1).c_str(),
                   kRows, kCols, (with_add ? "with add" : "without add"), i, exp,
                   act, tolerance, l1, rel);
           HWY_ASSERT(0);
         }
       }
     };

     assert_close(expected_without_add, actual_without_add, /*with_add=*/false);

     AlignedFreeUniquePtr<T[]> expected_with_add = AllocateAligned<T>(kRows);
     SimpleMatVecAdd(pm, pv, av, kRows, kCols, expected_with_add.get(), pool);

     AlignedFreeUniquePtr<T[]> actual_with_add = AllocateAligned<T>(kRows);
     MatVecAdd<kRows, kCols>(pm, pv, av, actual_with_add.get(), pool);

     assert_close(expected_with_add, actual_with_add, /*with_add=*/true);

 #else
     (void)d;
     (void)pool;
 #endif  // HWY_TARGET != HWY_SCALAR
   }

   template <class D>
   HWY_NOINLINE void CreatePoolAndTest(D d, size_t num_threads) {
     // Threads might not work on WASM; run only on main thread.
     if (HaveThreadingSupport()) num_threads = 0;

     ThreadPool pool(HWY_MIN(num_threads, ThreadPool::MaxThreads()));

     Test<AdjustedReps(192), AdjustedReps(256)>(d, pool);
 // Fewer tests due to compiler OOM
 #if !HWY_ARCH_RISCV
     Test<40, AdjustedReps(512)>(d, pool);
     Test<AdjustedReps(1024), 50>(d, pool);

     // Too large for low-precision vectors/accumulators.
     if (sizeof(TFromD<D>) != 2 && sizeof(VecT) != 2) {
       Test<AdjustedReps(1536), AdjustedReps(1536)>(d, pool);
     }
 #endif  // !HWY_ARCH_RISCV
   }

  public:
   template <class T, class D>
   HWY_NOINLINE void operator()(T /*unused*/, D d) {
     CreatePoolAndTest(d, 13);
 // Fewer tests due to compiler OOM
 #if !HWY_ARCH_RISCV
     CreatePoolAndTest(d, 16);
 #endif
   }
 };

 void TestAllMatVecAdd() {
 #if HWY_HAVE_FLOAT16
   ForPartialVectors<TestMatVecAdd<float16_t, float16_t>>()(float16_t());
 #endif
   ForPartialVectors<TestMatVecAdd<float, float>>()(float());
 #if HWY_HAVE_FLOAT64
   ForPartialVectors<TestMatVecAdd<double, double>>()(double());
 #endif
 }

 void TestAllMatVecBF16() {
   ForGEVectors<32, TestMatVecAdd<bfloat16_t, float>>()(float());
 }

 void TestAllMatVecBF16Both() {
   ForGEVectors<32, TestMatVecAdd<bfloat16_t, bfloat16_t>>()(float());
 }

 }  // namespace
 // NOLINTNEXTLINE(google-readability-namespace-comments)
 }  // namespace HWY_NAMESPACE
 }  // namespace hwy
 HWY_AFTER_NAMESPACE();

 #if HWY_ONCE
 namespace hwy {
 namespace {
 HWY_BEFORE_TEST(MatVecTest);
 HWY_EXPORT_AND_TEST_P(MatVecTest, TestAllMatVecAdd);
 HWY_EXPORT_AND_TEST_P(MatVecTest, TestAllMatVecBF16);
 HWY_EXPORT_AND_TEST_P(MatVecTest, TestAllMatVecBF16Both);
 HWY_AFTER_TEST();
 }  // namespace
 }  // namespace hwy
 HWY_TEST_MAIN();
 #endif  // HWY_ONCE
	// Copyright 2023 Google LLC
	// SPDX-License-Identifier: Apache-2.0
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	#include "third_party/highway/hwy/base.h"

	// Reduce targets to avoid timeout under emulation.
	#ifndef HWY_DISABLED_TARGETS
	#define HWY_DISABLED_TARGETS (HWY_SVE2_128 \| HWY_SVE2 \| HWY_SVE_256 \| HWY_NEON)
	#endif

	#include <stddef.h>
	#include <stdint.h>

	#include <cmath> // std::abs

	#include "third_party/highway/hwy/aligned_allocator.h"

	// clang-format off
	#undef HWY_TARGET_INCLUDE
	#define HWY_TARGET_INCLUDE "hwy/contrib/matvec/matvec_test.cc" // NOLINT
	#include "third_party/highway/hwy/foreach_target.h" // IWYU pragma: keep
	// Must come after foreach_target.h
	#include "third_party/highway/hwy/contrib/algo/transform-inl.h"
	#include "third_party/highway/hwy/contrib/matvec/matvec-inl.h"
	#include "third_party/highway/hwy/highway.h"
	#include "third_party/highway/hwy/contrib/thread_pool/thread_pool.h"
	#include "third_party/highway/hwy/contrib/thread_pool/topology.h"
	#include "third_party/highway/hwy/tests/test_util-inl.h"
	// clang-format on

	HWY_BEFORE_NAMESPACE();
	namespace hwy {
	namespace HWY_NAMESPACE {
	namespace {

	template <typename MatT, typename T>
	HWY_NOINLINE void SimpleMatVecAdd(const MatT* HWY_RESTRICT mat,
	const T* HWY_RESTRICT vec,
	const T* HWY_RESTRICT add, size_t rows,
	size_t cols, T* HWY_RESTRICT out,
	ThreadPool& pool) {
	if (add) {
	pool.Run(0, rows, [=](uint64_t r, size_t /thread/) {
	double dot = 0.0;
	for (size_t c = 0; c < cols; c++) {
	// For reasons unknown, fp16 += does not compile on clang (Arm).
	dot += ConvertScalarTo<double>(mat[r * cols + c]) *
	ConvertScalarTo<double>(vec[c]);
	}
	out[r] = ConvertScalarTo<T>(dot + ConvertScalarTo<double>(add[r]));
	});
	} else {
	pool.Run(0, rows, [=](uint64_t r, size_t /thread/) {
	double dot = 0.0;
	for (size_t c = 0; c < cols; c++) {
	// For reasons unknown, fp16 += does not compile on clang (Arm).
	dot += ConvertScalarTo<double>(mat[r * cols + c]) *
	ConvertScalarTo<double>(vec[c]);
	}
	out[r] = ConvertScalarTo<T>(dot);
	});
	}
	}

	HWY_MAYBE_UNUSED HWY_NOINLINE void SimpleMatVecAdd(
	const hwy::bfloat16_t* HWY_RESTRICT mat, const float* HWY_RESTRICT vec,
	const float* add, size_t rows, size_t cols, float* HWY_RESTRICT out,
	ThreadPool& pool) {
	if (add) {
	pool.Run(0, rows, [=](uint64_t r, size_t /thread/) {
	float dot = 0.0f;
	for (size_t c = 0; c < cols; c++) {
	dot += F32FromBF16(mat[r * cols + c]) * vec[c];
	}
	out[r] = dot + add[r];
	});
	} else {
	pool.Run(0, rows, [=](uint64_t r, size_t /thread/) {
	float dot = 0.0f;
	for (size_t c = 0; c < cols; c++) {
	dot += F32FromBF16(mat[r * cols + c]) * vec[c];
	}
	out[r] = dot;
	});
	}
	}

	HWY_MAYBE_UNUSED HWY_NOINLINE void SimpleMatVecAdd(
	const hwy::bfloat16_t* HWY_RESTRICT mat,
	const hwy::bfloat16_t* HWY_RESTRICT vec,
	const hwy::bfloat16_t* HWY_RESTRICT add, size_t rows, size_t cols,
	float* HWY_RESTRICT out, ThreadPool& pool) {
	if (add) {
	pool.Run(0, rows, [=](uint64_t r, size_t /thread/) {
	float dot = 0.0f;
	for (size_t c = 0; c < cols; c++) {
	dot += F32FromBF16(mat[r * cols + c]) * F32FromBF16(vec[c]);
	}
	out[r] = dot + F32FromBF16(add[r]);
	});
	} else {
	pool.Run(0, rows, [=](uint64_t r, size_t /thread/) {
	float dot = 0.0f;
	for (size_t c = 0; c < cols; c++) {
	dot += F32FromBF16(mat[r * cols + c]) * F32FromBF16(vec[c]);
	}
	out[r] = dot;
	});
	}
	}

	// Workaround for incorrect codegen on Arm, which results in values of `av`
	// >= 1E10. Can also be prevented by calling `Print(du, indices)`.
	#if HWY_ARCH_ARM && HWY_COMPILER_CLANG
	#define GENERATE_INLINE HWY_NOINLINE
	#else
	#define GENERATE_INLINE HWY_INLINE
	#endif

	struct GenerateMod {
	template <class D, HWY_IF_NOT_BF16_D(D), HWY_IF_LANES_GT_D(D, 1)>
	GENERATE_INLINE Vec<D> operator()(D d,
	Vec<RebindToUnsigned<D>> indices) const {
	const RebindToUnsigned<D> du;
	return Reverse2(d, ConvertTo(d, And(indices, Set(du, 0xF))));
	}

	template <class D, HWY_IF_NOT_BF16_D(D), HWY_IF_LANES_LE_D(D, 1)>
	GENERATE_INLINE Vec<D> operator()(D d,
	Vec<RebindToUnsigned<D>> indices) const {
	const RebindToUnsigned<D> du;
	return ConvertTo(d, And(indices, Set(du, 0xF)));
	}

	// Requires >= 4 bf16 lanes for float32 Reverse2.
	template <class D, HWY_IF_BF16_D(D), HWY_IF_LANES_GT_D(D, 2)>
	GENERATE_INLINE Vec<D> operator()(D d,
	Vec<RebindToUnsigned<D>> indices) const {
	const RebindToUnsigned<D> du;
	const RebindToSigned<D> di;
	const RepartitionToWide<decltype(di)> dw;
	const RebindToFloat<decltype(dw)> df;
	indices = And(indices, Set(du, 0xF));
	const Vec<decltype(df)> i0 = ConvertTo(df, PromoteLowerTo(dw, indices));
	const Vec<decltype(df)> i1 = ConvertTo(df, PromoteUpperTo(dw, indices));
	return OrderedDemote2To(d, Reverse2(df, i0), Reverse2(df, i1));
	}

	// For one or two lanes, we don't have OrderedDemote2To nor Reverse2.
	template <class D, HWY_IF_BF16_D(D), HWY_IF_LANES_LE_D(D, 2)>
	GENERATE_INLINE Vec<D> operator()(D d,
	Vec<RebindToUnsigned<D>> indices) const {
	const Rebind<float, D> df;
	return DemoteTo(d, Set(df, static_cast<float>(GetLane(indices))));
	}
	};

	// MatT is usually the same as T, but can also be bfloat16_t when T = float.
	template <typename MatT, typename VecT>
	class TestMatVecAdd {
	template <size_t kRows, size_t kCols, class D, typename T = TFromD<D>>
	HWY_NOINLINE void Test(D d, ThreadPool& pool) {
	// This target lacks too many ops required in our implementation, use
	// HWY_EMU128 instead.
	#if HWY_TARGET != HWY_SCALAR
	const Repartition<MatT, D> dm;
	const Repartition<VecT, D> dv;
	const size_t misalign = 3 * Lanes(d) / 5;
	// Fill matrix and vector with small integer values
	const size_t area = kRows * kCols;
	AlignedFreeUniquePtr<MatT[]> storage_m =
	AllocateAligned<MatT>(misalign + area);
	AlignedFreeUniquePtr<VecT[]> storage_v =
	AllocateAligned<VecT>(misalign + kCols);
	AlignedFreeUniquePtr<VecT[]> storage_a =
	AllocateAligned<VecT>(misalign + kRows);
	HWY_ASSERT(storage_m && storage_v && storage_a);
	MatT* pm = storage_m.get() + misalign;
	VecT* pv = storage_v.get() + misalign;
	VecT* av = storage_a.get() + misalign;
	Generate(dm, pm, area, GenerateMod());
	Generate(dv, pv, kCols, GenerateMod());
	Generate(dv, av, kRows, GenerateMod());

	AlignedFreeUniquePtr<T[]> expected_without_add = AllocateAligned<T>(kRows);
	HWY_ASSERT(expected_without_add);
	SimpleMatVecAdd(pm, pv, static_cast<VecT*>(nullptr), kRows, kCols,
	expected_without_add.get(), pool);

	AlignedFreeUniquePtr<T[]> actual_without_add = AllocateAligned<T>(kRows);
	HWY_ASSERT(actual_without_add);
	MatVec<kRows, kCols>(pm, pv, actual_without_add.get(), pool);

	const auto assert_close = [&](const AlignedFreeUniquePtr<T[]>& expected,
	const AlignedFreeUniquePtr<T[]>& actual,
	bool with_add) {
	for (size_t i = 0; i < kRows; ++i) {
	const double exp = ConvertScalarTo<double>(expected[i]);
	const double act = ConvertScalarTo<double>(actual[i]);
	const double epsilon =
	1.0 / (1ULL << HWY_MIN(MantissaBits<MatT>(), MantissaBits<VecT>()));
	const double tolerance = exp * 20.0 / epsilon;
	const double l1 = std::abs(exp - act);
	const double rel = exp == 0.0 ? 0.0 : l1 / exp;

	if (l1 > tolerance && rel > epsilon) {
	fprintf(stderr,
	"%s/%s %zu x %zu, %s: mismatch at %zu: %E != %E; "
	"tol %f l1 %f rel %E\n",
	TypeName(MatT(), 1).c_str(), TypeName(VecT(), 1).c_str(),
	kRows, kCols, (with_add ? "with add" : "without add"), i, exp,
	act, tolerance, l1, rel);
	HWY_ASSERT(0);
	}
	}
	};

	assert_close(expected_without_add, actual_without_add, /with_add=/false);

	AlignedFreeUniquePtr<T[]> expected_with_add = AllocateAligned<T>(kRows);
	SimpleMatVecAdd(pm, pv, av, kRows, kCols, expected_with_add.get(), pool);

	AlignedFreeUniquePtr<T[]> actual_with_add = AllocateAligned<T>(kRows);
	MatVecAdd<kRows, kCols>(pm, pv, av, actual_with_add.get(), pool);

	assert_close(expected_with_add, actual_with_add, /with_add=/true);

	#else
	(void)d;
	(void)pool;
	#endif // HWY_TARGET != HWY_SCALAR
	}

	template <class D>
	HWY_NOINLINE void CreatePoolAndTest(D d, size_t num_threads) {
	// Threads might not work on WASM; run only on main thread.
	if (HaveThreadingSupport()) num_threads = 0;

	ThreadPool pool(HWY_MIN(num_threads, ThreadPool::MaxThreads()));

	Test<AdjustedReps(192), AdjustedReps(256)>(d, pool);
	// Fewer tests due to compiler OOM
	#if !HWY_ARCH_RISCV
	Test<40, AdjustedReps(512)>(d, pool);
	Test<AdjustedReps(1024), 50>(d, pool);

	// Too large for low-precision vectors/accumulators.
	if (sizeof(TFromD<D>) != 2 && sizeof(VecT) != 2) {
	Test<AdjustedReps(1536), AdjustedReps(1536)>(d, pool);
	}
	#endif // !HWY_ARCH_RISCV
	}

	public:
	template <class T, class D>
	HWY_NOINLINE void operator()(T /unused/, D d) {
	CreatePoolAndTest(d, 13);
	// Fewer tests due to compiler OOM
	#if !HWY_ARCH_RISCV
	CreatePoolAndTest(d, 16);
	#endif
	}
	};

	void TestAllMatVecAdd() {
	#if HWY_HAVE_FLOAT16
	ForPartialVectors<TestMatVecAdd<float16_t, float16_t>>()(float16_t());
	#endif
	ForPartialVectors<TestMatVecAdd<float, float>>()(float());
	#if HWY_HAVE_FLOAT64
	ForPartialVectors<TestMatVecAdd<double, double>>()(double());
	#endif
	}

	void TestAllMatVecBF16() {
	ForGEVectors<32, TestMatVecAdd<bfloat16_t, float>>()(float());
	}

	void TestAllMatVecBF16Both() {
	ForGEVectors<32, TestMatVecAdd<bfloat16_t, bfloat16_t>>()(float());
	}

	} // namespace
	// NOLINTNEXTLINE(google-readability-namespace-comments)
	} // namespace HWY_NAMESPACE
	} // namespace hwy
	HWY_AFTER_NAMESPACE();

	#if HWY_ONCE
	namespace hwy {
	namespace {
	HWY_BEFORE_TEST(MatVecTest);
	HWY_EXPORT_AND_TEST_P(MatVecTest, TestAllMatVecAdd);
	HWY_EXPORT_AND_TEST_P(MatVecTest, TestAllMatVecBF16);
	HWY_EXPORT_AND_TEST_P(MatVecTest, TestAllMatVecBF16Both);
	HWY_AFTER_TEST();
	} // namespace
	} // namespace hwy
	HWY_TEST_MAIN();
	#endif // HWY_ONCE