Clean up SIMD headers and tests

- Be more precise/specific about the mapping from architecture to header
- Add more test coverage by giving up on using gtest to instantiate the tests
- Fix bugs that motivated these changes: AVX512 s32 multiply was wrong, sse41 intrinsics were enabled from sse2 headers
- Refactor arm_neon_xf16.h to be unnecessary, multi_vec.h can do this `extract` now.

PiperOrigin-RevId: 840525070

Author: dsharletg

ynnpack/base/simd/BUILD

@@ -18,11 +18,12 @@ cc_library(
         "//ynnpack:x86": [
             "x86_avx.h",
             "x86_avx2.h",
-            "x86_avx512.h",
-            "x86_sse.h",
+            "x86_avx512bw.h",
+            "x86_avx512f.h",
+            "x86_sse2.h",
             "x86_sse41.h",
         ],
-        "//ynnpack:arm": ["arm.h"],
+        "//ynnpack:arm": ["arm_neon.h"],
         "//conditions:default": [],
     }),
     compatible_with = _COMPATIBLE_WITH,

ynnpack/base/simd/arm.h renamed to ynnpack/base/simd/arm_neon.h

@@ -124,6 +124,12 @@ YNN_ALWAYS_INLINE void store(typename Vec::value_type* ptr,
   }
 }
 
+template <int Index, typename Vec>
+Vec extract(Vec x, Vec) {
+  static_assert(Index == 0, "");
+  return x;
+}
+
 template <int Index, typename Vec, typename ResultVec, size_t M>
 ResultVec extract(multi_vec<Vec, M> b, ResultVec) {
   static_assert(Index * ResultVec::N < M * Vec::N);

ynnpack/base/simd/multi_vec.h

@@ -11,27 +11,16 @@ package(default_visibility = ["//ynnpack:__subpackages__"])
 ynn_cc_library(
     name = "tests",
     testonly = True,
+    hdrs = ["generic.h"],
     per_arch_srcs = {
         # This test defines different tests depending on which architecture is being compiled.
-        "x86_sse2": [
-            "test.cc",
-            "convert_only_test.cc",
-        ],
-        "x86_sse41": ["test.cc"],
-        "x86_avx": ["test.cc"],
-        "x86_avx2": [
-            "test.cc",
-            "convert_only_test.cc",
-        ],
-        "x86_f16c": ["convert_only_test.cc"],
-        "x86_avx512bf16": ["convert_only_test.cc"],
-        "x86_avx512fp16": ["convert_only_test.cc"],
-        "x86_avx512f": ["test.cc"],
-        "x86_avx512bw": ["test.cc"],
-        "arm_neon": [
-            "convert_only_test.cc",
-            "test.cc",
-        ],
+        "x86_sse2": ["x86_sse2.cc"],
+        "x86_sse41": ["x86_sse41.cc"],
+        "x86_avx": ["x86_avx.cc"],
+        "x86_avx2": ["x86_avx2.cc"],
+        "x86_avx512f": ["x86_avx512f.cc"],
+        "x86_avx512bw": ["x86_avx512bw.cc"],
+        "arm_neon": ["arm_neon.cc"],
     },
     deps = [
         "//ynnpack/base",  # buildcleaner:keep

ynnpack/base/simd/test/BUILD

@@ -0,0 +1,77 @@
+// Copyright 2025 Google LLC
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "ynnpack/base/simd/arm_neon.h"
+
+#include "ynnpack/base/simd/test/generic.h"
+
+namespace ynn {
+namespace simd {
+
+TEST_LOAD_STORE(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_LOAD_STORE(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_LOAD_STORE(arm_neon, s16x8, /*arch_flags=*/0);
+TEST_LOAD_STORE(arm_neon, f16x8, /*arch_flags=*/0);
+TEST_LOAD_STORE(arm_neon, bf16x8, /*arch_flags=*/0);
+TEST_LOAD_STORE(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_LOAD_STORE(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_ALIGNED_LOAD_STORE(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_ALIGNED_LOAD_STORE(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_ALIGNED_LOAD_STORE(arm_neon, s16x8, /*arch_flags=*/0);
+TEST_ALIGNED_LOAD_STORE(arm_neon, f16x8, /*arch_flags=*/0);
+TEST_ALIGNED_LOAD_STORE(arm_neon, bf16x8, /*arch_flags=*/0);
+TEST_ALIGNED_LOAD_STORE(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_ALIGNED_LOAD_STORE(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_PARTIAL_LOAD_STORE(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_PARTIAL_LOAD_STORE(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_PARTIAL_LOAD_STORE(arm_neon, s16x8, /*arch_flags=*/0);
+TEST_PARTIAL_LOAD_STORE(arm_neon, f16x8, /*arch_flags=*/0);
+TEST_PARTIAL_LOAD_STORE(arm_neon, bf16x8, /*arch_flags=*/0);
+TEST_PARTIAL_LOAD_STORE(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_PARTIAL_LOAD_STORE(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_ADD(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_ADD(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_ADD(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_ADD(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_SUBTRACT(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_SUBTRACT(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_SUBTRACT(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_SUBTRACT(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_MULTIPLY(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_MULTIPLY(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_MULTIPLY(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_MULTIPLY(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_MIN(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_MIN(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_MIN(arm_neon, s16x8, /*arch_flags=*/0);
+TEST_MIN(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_MIN(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_MAX(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_MAX(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_MAX(arm_neon, s16x8, /*arch_flags=*/0);
+TEST_MAX(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_MAX(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_HORIZONTAL_MIN(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_HORIZONTAL_MIN(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_HORIZONTAL_MIN(arm_neon, s16x8, /*arch_flags=*/0);
+TEST_HORIZONTAL_MIN(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_HORIZONTAL_MIN(arm_neon, s32x4, /*arch_flags=*/0);
+
+TEST_HORIZONTAL_MAX(arm_neon, u8x16, /*arch_flags=*/0);
+TEST_HORIZONTAL_MAX(arm_neon, s8x16, /*arch_flags=*/0);
+TEST_HORIZONTAL_MAX(arm_neon, s16x8, /*arch_flags=*/0);
+TEST_HORIZONTAL_MAX(arm_neon, f32x4, /*arch_flags=*/0);
+TEST_HORIZONTAL_MAX(arm_neon, s32x4, /*arch_flags=*/0);
+
+}  // namespace simd
+}  // namespace ynn