diff --git a/hwy/ops/arm_neon-inl.h b/hwy/ops/arm_neon-inl.h index 2b8251076e..3dc88c05e6 100644 --- a/hwy/ops/arm_neon-inl.h +++ b/hwy/ops/arm_neon-inl.h @@ -7834,8 +7834,8 @@ HWY_API VFromD SumOfMulQuadAccumulate( #define HWY_NATIVE_U8_I8_SUMOFMULQUADACCUMULATE #endif -#if defined(__ARM_FEATURE_MATMUL_INT8) || \ - (HWY_TARGET == HWY_NEON_BF16 && HWY_OS_APPLE && HWY_ARCH_ARM_A64 && \ +#if defined(__ARM_FEATURE_MATMUL_INT8) || \ + (HWY_TARGET == HWY_NEON_BF16 && HWY_OS_APPLE && \ HWY_HAVE_RUNTIME_DISPATCH) template @@ -10958,6 +10958,31 @@ namespace detail { // for code folding #undef HWY_NEON_IF_NOT_EMULATED_D } // namespace detail +// ------------------------------ PerBlock2x2MatMul (Neon hardware overrides) +#if defined(__ARM_FEATURE_MATMUL_INT8) || \ + (HWY_TARGET == HWY_NEON_BF16 && HWY_OS_APPLE && \ + HWY_HAVE_RUNTIME_DISPATCH) +template +HWY_API Vec128 PerBlock2x2MatMul( + Simd /* d */, + Vec128 a, + Vec128 b, + Vec128 c) { + return Vec128{vmmlaq_s32(c.raw, a.raw, b.raw)}; +} +#endif + +#if defined(__ARM_FEATURE_BF16_VECTOR_ARITHMETIC) || HWY_TARGET == HWY_NEON_BF16 +template +HWY_API Vec128 PerBlock2x2MatMul( + Simd /* d */, + Vec128 a, + Vec128 b, + Vec128 c) { + return Vec128{vbfmmlaq_f32(c.raw, a.raw, b.raw)}; +} +#endif + // NOLINTNEXTLINE(google-readability-namespace-comments) } // namespace HWY_NAMESPACE } // namespace hwy diff --git a/hwy/ops/arm_sve-inl.h b/hwy/ops/arm_sve-inl.h index c90bf19dde..950754ea6c 100644 --- a/hwy/ops/arm_sve-inl.h +++ b/hwy/ops/arm_sve-inl.h @@ -7109,6 +7109,33 @@ HWY_SVE_FOREACH_UI(HWY_SVE_MASKED_LEADING_ZERO_COUNT, MaskedLeadingZeroCount, clz) #undef HWY_SVE_LEADING_ZERO_COUNT +// ------------------------------ PerBlock2x2MatMul (SVE hardware overrides) +#if defined(__ARM_FEATURE_SVE_MATMUL_INT8) || HWY_TARGET == HWY_SVE2_128 +template +HWY_API svint32_t PerBlock2x2MatMul( + Simd /* d */, + svint8_t a, + svint8_t b, + svint32_t c) { + return svmmla_s32(c, a, b); +} +#endif + +#if defined(__ARM_FEATURE_SVE_BF16) || HWY_TARGET == HWY_SVE2_128 +template +HWY_API svfloat32_t PerBlock2x2MatMul( + Simd /* d */, + VBF16 a, + VBF16 b, + svfloat32_t c) { +#if HWY_SVE_HAVE_BF16_VEC + return svbfmmla_f32(c, a, b); +#else + return svbfmmla_f32(c, svreinterpret_bf16_u16(a), svreinterpret_bf16_u16(b)); +#endif +} +#endif + // ================================================== END MACROS #undef HWY_SVE_ALL_PTRUE #undef HWY_SVE_D diff --git a/hwy/ops/set_macros-inl.h b/hwy/ops/set_macros-inl.h index 4fe0e9d35f..4db9e9ab0b 100644 --- a/hwy/ops/set_macros-inl.h +++ b/hwy/ops/set_macros-inl.h @@ -49,6 +49,11 @@ #undef HWY_NATIVE_MASK #undef HWY_NATIVE_INTERLEAVE_WHOLE +#undef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 +#undef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 0 +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 0 + #ifndef HWY_CAP_GE256 #define HWY_CAP_GE256 (HWY_MIN_BYTES >= 32) #endif @@ -530,6 +535,22 @@ #define HWY_NATIVE_DOT_BF16 0 #endif +#undef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 +#if defined(__ARM_FEATURE_MATMUL_INT8) || \ + (HWY_TARGET == HWY_NEON_BF16 && HWY_OS_APPLE && \ + HWY_HAVE_RUNTIME_DISPATCH) +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 1 +#else +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 0 +#endif + +#undef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 +#if defined(__ARM_FEATURE_BF16_VECTOR_ARITHMETIC) || HWY_TARGET == HWY_NEON_BF16 +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 1 +#else +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 0 +#endif + #define HWY_NATIVE_MASK 0 #if HWY_TARGET == HWY_NEON_WITHOUT_AES @@ -621,6 +642,19 @@ #else #define HWY_NATIVE_DOT_BF16 0 #endif +#undef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 +#if defined(__ARM_FEATURE_SVE_MATMUL_INT8) || HWY_TARGET == HWY_SVE2_128 +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 1 +#else +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 0 +#endif + +#undef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 +#if defined(__ARM_FEATURE_SVE_BF16) || HWY_TARGET == HWY_SVE2_128 +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 1 +#else +#define HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 0 +#endif #define HWY_NATIVE_MASK 1 #if HWY_TARGET == HWY_SVE2 diff --git a/hwy/tests/matmul_op_test.cc b/hwy/tests/matmul_op_test.cc new file mode 100644 index 0000000000..a739463151 --- /dev/null +++ b/hwy/tests/matmul_op_test.cc @@ -0,0 +1,164 @@ +// Copyright 2026 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. + +#undef HWY_TARGET_INCLUDE +#define HWY_TARGET_INCLUDE "tests/matmul_op_test.cc" +#include "hwy/foreach_target.h" // IWYU pragma: keep +#include "hwy/highway.h" +#include "hwy/tests/test_util-inl.h" + +HWY_BEFORE_NAMESPACE(); +namespace hwy { +namespace HWY_NAMESPACE { +namespace { + +#ifndef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 +#error "Bug in set_macros-inl.h, did not set HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8" +#endif + +#ifndef HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 +#error "Bug in set_macros-inl.h, did not set HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16" +#endif + +struct TestInt8PerBlock2x2MatMul { + template + HWY_NOINLINE void operator()(TN /*unused*/, DN dn) { +#if HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 + static_assert(IsSame(), "TN should be int32_t"); + const Repartition di8; + using VI8 = Vec; + using V32 = Vec; + HWY_LANES_CONSTEXPR size_t N = Lanes(dn); + + // Allocate aligned buffers for scalar verification + auto in_a = AllocateAligned(N * 4); + auto in_b = AllocateAligned(N * 4); + auto in_c = AllocateAligned(N); + auto expected = AllocateAligned(N); + HWY_ASSERT(in_a && in_b && in_c && expected); + + // Populate buffers + for (size_t i = 0; i < N * 4; ++i) { + in_a[i] = static_cast((i % 25) - 12); + in_b[i] = static_cast((i % 19) - 9); + } + for (size_t i = 0; i < N; ++i) { + in_c[i] = static_cast((i % 7) + 10); + expected[i] = in_c[i]; + } + + // Scalar emulation loop (matching hardware svmmla interleaving) + for (size_t block = 0; block < N; block += 4) { + const size_t block_i8 = block * 4; + for (int i = 0; i < 2; ++i) { + for (int j = 0; j < 2; ++j) { + int32_t sum = 0; + for (int k = 0; k < 8; ++k) { + sum += static_cast(in_a[block_i8 + i * 8 + k]) * + static_cast(in_b[block_i8 + j * 8 + k]); + } + expected[block + i * 2 + j] += sum; + } + } + } + + const VI8 va = Load(di8, in_a.get()); + const VI8 vb = Load(di8, in_b.get()); + const V32 vc = Load(dn, in_c.get()); + + const V32 actual = PerBlock2x2MatMul(dn, va, vb, vc); + HWY_ASSERT_VEC_EQ(dn, expected.get(), actual); +#else + (void)dn; +#endif + } +}; + +HWY_NOINLINE void TestAllInt8PerBlock2x2MatMul() { + ForGEVectors<128, TestInt8PerBlock2x2MatMul>()(int32_t()); +} + +struct TestBf16PerBlock2x2MatMul { + template + HWY_NOINLINE void operator()(TN /*unused*/, DN dn) { +#if HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16 + static_assert(IsSame(), "TN should be float"); + const Repartition dbf; + using VBF = Vec; + using VF = Vec; + HWY_LANES_CONSTEXPR size_t N = Lanes(dn); + + auto in_a = AllocateAligned(N * 2); + auto in_b = AllocateAligned(N * 2); + auto in_c = AllocateAligned(N); + auto expected = AllocateAligned(N); + HWY_ASSERT(in_a && in_b && in_c && expected); + + for (size_t i = 0; i < N * 2; ++i) { + in_a[i] = hwy::ConvertScalarTo(static_cast(i % 5) * 0.5f); + in_b[i] = hwy::ConvertScalarTo(static_cast(i % 7) * 0.25f); + } + for (size_t i = 0; i < N; ++i) { + in_c[i] = static_cast(i % 3) + 1.0f; + expected[i] = in_c[i]; + } + + for (size_t block = 0; block < N; block += 4) { + const size_t block_bf = block * 2; + for (int i = 0; i < 2; ++i) { + for (int j = 0; j < 2; ++j) { + float sum = 0.0f; + for (int k = 0; k < 4; ++k) { + sum += hwy::ConvertScalarTo(in_a[block_bf + i * 4 + k]) * + hwy::ConvertScalarTo(in_b[block_bf + j * 4 + k]); + } + expected[block + i * 2 + j] += sum; + } + } + } + + const VBF va = Load(dbf, in_a.get()); + const VBF vb = Load(dbf, in_b.get()); + const VF vc = Load(dn, in_c.get()); + + const VF actual = PerBlock2x2MatMul(dn, va, vb, vc); + HWY_ASSERT_VEC_EQ(dn, expected.get(), actual); +#else + (void)dn; +#endif + } +}; + +HWY_NOINLINE void TestAllBf16PerBlock2x2MatMul() { + ForGEVectors<128, TestBf16PerBlock2x2MatMul>()(float()); +} + +} // namespace +// NOLINTNEXTLINE(google-readability-namespace-comments) +} // namespace HWY_NAMESPACE +} // namespace hwy +HWY_AFTER_NAMESPACE(); + +#if HWY_ONCE +namespace hwy { +namespace { +HWY_BEFORE_TEST(HwyMatmulOpTest); +HWY_EXPORT_AND_TEST_P(HwyMatmulOpTest, TestAllInt8PerBlock2x2MatMul); +HWY_EXPORT_AND_TEST_P(HwyMatmulOpTest, TestAllBf16PerBlock2x2MatMul); +HWY_AFTER_TEST(); +} // namespace +} // namespace hwy +HWY_TEST_MAIN(); +#endif // HWY_ONCE diff --git a/hwy_tests.bzl b/hwy_tests.bzl index feaca9e313..8a3eefb48d 100644 --- a/hwy_tests.bzl +++ b/hwy_tests.bzl @@ -232,6 +232,7 @@ HWY_TESTS = HWY_CONTRIB_TESTS + ( ("hwy/tests/", "masked_arithmetic_test", []), ("hwy/tests/", "masked_compare_test", []), ("hwy/tests/", "masked_minmax_test", []), + ("hwy/tests/", "matmul_op_test", []), ("hwy/tests/", "memory_test", []), ("hwy/tests/", "minmax_magnitude_test", []), ("hwy/tests/", "minmax_number_test", []),