Helper functions to use arm matrix multiplication instructions in a way that is testable on other platofrms

hwy/ops/arm_neon-inl.h

@@ -7834,8 +7834,8 @@ HWY_API VFromD<DU32> 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 <class DI32, HWY_IF_I32_D(DI32), HWY_IF_V_SIZE_LE_D(DI32, 8)>
@@ -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 <size_t N>
+HWY_API Vec128<int32_t, N> PerBlock2x2MatMul(
+    Simd<int32_t, N, 0> /* d */,
+    Vec128<int8_t, N * 4> a,
+    Vec128<int8_t, N * 4> b,
+    Vec128<int32_t, N> c) {
+  return Vec128<int32_t, N>{vmmlaq_s32(c.raw, a.raw, b.raw)};
+}
+#endif
+
+#if defined(__ARM_FEATURE_BF16_VECTOR_ARITHMETIC) || HWY_TARGET == HWY_NEON_BF16
+template <size_t N>
+HWY_API Vec128<float, N> PerBlock2x2MatMul(
+    Simd<float, N, 0> /* d */,
+    Vec128<hwy::bfloat16_t, N * 2> a,
+    Vec128<hwy::bfloat16_t, N * 2> b,
+    Vec128<float, N> c) {
+  return Vec128<float, N>{vbfmmlaq_f32(c.raw, a.raw, b.raw)};
+}
+#endif
+
 // NOLINTNEXTLINE(google-readability-namespace-comments)
 }  // namespace HWY_NAMESPACE
 }  // namespace hwy

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 <size_t N, int kPow2>
+HWY_API svint32_t PerBlock2x2MatMul(
+    Simd<int32_t, N, kPow2> /* 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 <size_t N, int kPow2, class VBF16>
+HWY_API svfloat32_t PerBlock2x2MatMul(
+    Simd<float, N, kPow2> /* 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

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

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 <typename TN, class DN>
+  HWY_NOINLINE void operator()(TN /*unused*/, DN dn) {
+#if HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8
+    static_assert(IsSame<TN, int32_t>(), "TN should be int32_t");
+    const Repartition<int8_t, DN> di8;
+    using VI8 = Vec<decltype(di8)>;
+    using V32 = Vec<decltype(dn)>;
+    HWY_LANES_CONSTEXPR size_t N = Lanes(dn);
+
+    // Allocate aligned buffers for scalar verification
+    auto in_a = AllocateAligned<int8_t>(N * 4);
+    auto in_b = AllocateAligned<int8_t>(N * 4);
+    auto in_c = AllocateAligned<int32_t>(N);
+    auto expected = AllocateAligned<int32_t>(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<int8_t>((i % 25) - 12);
+      in_b[i] = static_cast<int8_t>((i % 19) - 9);
+    }
+    for (size_t i = 0; i < N; ++i) {
+      in_c[i] = static_cast<int32_t>((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<int32_t>(in_a[block_i8 + i * 8 + k]) *
+                   static_cast<int32_t>(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 <typename TN, class DN>
+  HWY_NOINLINE void operator()(TN /*unused*/, DN dn) {
+#if HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16
+    static_assert(IsSame<TN, float>(), "TN should be float");
+    const Repartition<hwy::bfloat16_t, DN> dbf;
+    using VBF = Vec<decltype(dbf)>;
+    using VF = Vec<decltype(dn)>;
+    HWY_LANES_CONSTEXPR size_t N = Lanes(dn);
+
+    auto in_a = AllocateAligned<hwy::bfloat16_t>(N * 2);
+    auto in_b = AllocateAligned<hwy::bfloat16_t>(N * 2);
+    auto in_c = AllocateAligned<float>(N);
+    auto expected = AllocateAligned<float>(N);
+    HWY_ASSERT(in_a && in_b && in_c && expected);
+
+    for (size_t i = 0; i < N * 2; ++i) {
+      in_a[i] = hwy::ConvertScalarTo<hwy::bfloat16_t>(static_cast<float>(i % 5) * 0.5f);
+      in_b[i] = hwy::ConvertScalarTo<hwy::bfloat16_t>(static_cast<float>(i % 7) * 0.25f);
+    }
+    for (size_t i = 0; i < N; ++i) {
+      in_c[i] = static_cast<float>(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<float>(in_a[block_bf + i * 4 + k]) *
+                   hwy::ConvertScalarTo<float>(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

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", []),