feat: replace argmaxpool yaml with table header changes

BUILD.bazel

@@ -102,6 +102,8 @@ MICROKERNEL_DEFS = [
     "src/f16-vsqrt/f16-vsqrt.h",
     "src/f16-vtanh/f16-vtanh.h",
     "src/f32-avgpool/f32-avgpool-minmax.h",
+    "src/f32-argmaxpool/f32-argmaxpool-unipass.h",
+    "src/f32-argmaxpool/f32-argmaxpool-multipass.h",
     "src/f32-dwconv/f32-dwconv-minmax-multipass.h",
     "src/f32-dwconv/f32-dwconv-minmax-unipass.h",
     "src/f32-dwconv/f32-dwconv-multipass.h",

CMakeLists.txt

@@ -1265,7 +1265,6 @@ IF(XNNPACK_BUILD_TESTS)
     f32-vlrelu
     f32-rdsum
     f32-velu
-    f32-argmaxpool
     s8-vclamp
     u8-vclamp
   )
@@ -1286,6 +1285,28 @@ IF(XNNPACK_BUILD_TESTS)
     ADD_SHARDED_TEST(${TEST}-test 10)
   ENDFOREACH()
 
+SET(MICROKERNEL_ARGMAXPOOL_UNIT_TESTS
+      f32-argmaxpool-unipass
+      f32-argmaxpool-multipass)
+  FOREACH(TEST ${MICROKERNEL_ARGMAXPOOL_UNIT_TESTS})
+    ADD_EXECUTABLE(${TEST}-test test/${TEST}.cc)
+    TARGET_INCLUDE_DIRECTORIES(${TEST}-test PRIVATE include src test)
+    TARGET_LINK_LIBRARIES(${TEST}-test PRIVATE
+        GTest::gmock
+        GTest::gtest
+        GTest::gtest_main
+        hardware-config
+        logging
+        microkernels-all
+        microparams-init
+        next-prime
+        pthreadpool
+    )
+    ADD_TEST(NAME ${TEST}-test COMMAND ${TEST}-test)
+    SET_TARGET_PROPERTIES(${TEST}-test PROPERTIES CXX_EXTENSIONS YES)
+  ENDFOREACH()
+
+
   IF(XNNPACK_BUILD_LIBRARY)
     # ---[ Launch heavy tests first.
     SET(LIBRARY_SHARDED_TESTS

scripts/generate-tests.sh

@@ -5,7 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 ### Tests for ArgMaxPool micro-kernels
-tools/generate-argmaxpool-test.py --spec test/f32-argmaxpool.yaml --output test/f32-argmaxpool.cc &
+tools/generate-argmaxpool-unipass-test.py --tester ArgMaxPoolMicrokernelTester --ukernel f32-argmaxpool-unipass --output test/f32-argmaxpool-unipass.cc &
+tools/generate-argmaxpool-multipass-test.py --tester ArgMaxPoolMicrokernelTester --ukernel f32-argmaxpool-multipass --output test/f32-argmaxpool-multipass.cc &
 
 ### Tests for GEMM micro-kernels
 tools/generate-gemm-test.py --spec test/bf16-gemm-minmax.yaml --output-test test/bf16-gemm-minmax.cc &

src/f32-argmaxpool/f32-argmaxpool-multipass.h

@@ -0,0 +1,50 @@
+// Copyright 2024 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.
+
+#ifndef XNN_UKERNEL_WITH_PARAMS
+#define XNN_UKERNEL_WITH_PARAMS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype, params_type, init_params) \
+    XNN_UKERNEL_MULTIPASS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype)
+#define XNN_DEFINED_UKERNEL_WITH_PARAMS
+#endif
+
+#ifndef XNN_UKERNEL_MULTIPASS
+#define XNN_UKERNEL_MULTIPASS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype) \
+    XNN_UKERNEL_WITH_PARAMS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype, void, /*init_params=*/nullptr)
+#define XNN_DEFINED_UKERNEL
+#endif
+
+//SCALAR
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_9p8x__scalar_c1, 9, 8, 1, false, float, struct xnn_f32_default_params, NULL)
+
+
+#if (XNN_ARCH_ARM || XNN_ARCH_ARM64)
+XNN_UKERNEL_WITH_PARAMS(xnn_arch_arm_neon, xnn_f32_argmaxpool_ukernel_9p8x__neon_c4, 9, 8, 4,false, float, struct xnn_f32_default_params, NULL)
+
+
+#endif  // XNN_ARCH_ARM || XNN_ARCH_ARM64
+
+#if XNN_ARCH_X86 || XNN_ARCH_X86_64
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_9p8x__sse2_c4, 9, 8, 4, false, float, struct xnn_f32_default_params, NULL)
+
+#endif  // XNN_ARCH_X86 || XNN_ARCH_X86_64
+
+#if XNN_ARCH_WASMSIMD || XNN_ARCH_WASMRELAXEDSIMD
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_9p8x__wasmsimd_c4, 9, 8, 4, false, float, struct xnn_f32_default_params, NULL)
+#endif  // XNN_ARCH_WASMSIMD || XNN_ARCH_WASMRELAXEDSIMD
+
+#if XNN_ARCH_RISCV && (XNN_ENABLE_RISCV_VECTOR)
+
+XNN_UKERNEL_WITH_PARAMS(xnn_arch_riscv_vector, xnn_f32_argmaxpool_ukernel_9p8x__rvv_u1v,  9, 8, 1, true, float, struct xnn_f32_default_params, NULL)
+#endif // XNN_ARCH_RISCV && XNN_ENABLE_RISCV_VECTOR
+
+#ifdef XNN_DEFINED_UKERNEL_WITH_PARAMS
+#undef XNN_DEFINED_UKERNEL_WITH_PARAMS
+#undef XNN_UKERNEL_WITH_PARAMS
+#endif
+
+#ifdef XNN_DEFINED_UKERNEL
+#undef XNN_DEFINED_UKERNEL
+#undef XNN_UKERNEL_MULTIPASS
+#endif

src/f32-argmaxpool/f32-argmaxpool-unipass.h

@@ -0,0 +1,53 @@
+// Copyright 2024 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.
+
+#ifndef XNN_UKERNEL_WITH_PARAMS
+#define XNN_UKERNEL_WITH_PARAMS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype, params_type, init_params) \
+    XNN_UKERNEL_UNIPASS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype)
+#define XNN_DEFINED_UKERNEL_WITH_PARAMS
+#endif
+
+#ifndef XNN_UKERNEL_UNIPASS
+#define XNN_UKERNEL_UNIPASS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype) \
+    XNN_UKERNEL_WITH_PARAMS(arch_flags, ukernel, primary_tile, incremental_tile, channel_tile, vector_tile, datatype, void, /*init_params=*/nullptr)
+#define XNN_DEFINED_UKERNEL
+#endif
+
+//SCALAR
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_4x__scalar_c1, 4, 0, 1, false, float, struct xnn_f32_default_params, NULL)
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_9x__scalar_c1, 9, 0, 1, false, float, struct xnn_f32_default_params, NULL)
+
+#if (XNN_ARCH_ARM || XNN_ARCH_ARM64)
+XNN_UKERNEL_WITH_PARAMS(xnn_arch_arm_neon, xnn_f32_argmaxpool_ukernel_4x__neon_c4, 4, 0, 4,false, float, struct xnn_f32_default_params, NULL)
+XNN_UKERNEL_WITH_PARAMS(xnn_arch_arm_neon, xnn_f32_argmaxpool_ukernel_9x__neon_c4, 9, 0, 4, false,float, struct xnn_f32_default_params, NULL)
+
+#endif  // XNN_ARCH_ARM || XNN_ARCH_ARM64
+
+#if XNN_ARCH_X86 || XNN_ARCH_X86_64
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_4x__sse2_c4, 4, 0, 4, false, float, struct xnn_f32_default_params, NULL)
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_9x__sse2_c4,  9, 0, 4, false, float, struct xnn_f32_default_params, NULL)
+#endif  // XNN_ARCH_X86 || XNN_ARCH_X86_64
+
+#if XNN_ARCH_WASMSIMD || XNN_ARCH_WASMRELAXEDSIMD
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_4x__wasmsimd_c4, 4, 0, 4, false, float, struct xnn_f32_default_params, NULL)
+XNN_UKERNEL_WITH_PARAMS(0, xnn_f32_argmaxpool_ukernel_9x__wasmsimd_c4, 9, 0, 4, false, float, struct xnn_f32_default_params, NULL)
+#endif  // XNN_ARCH_WASMSIMD || XNN_ARCH_WASMRELAXEDSIMD
+
+#if XNN_ARCH_RISCV && (XNN_ENABLE_RISCV_VECTOR)
+
+XNN_UKERNEL_WITH_PARAMS(xnn_arch_riscv_vector, xnn_f32_argmaxpool_ukernel_4x__rvv_u1v,  4, 0, 1, true, float, struct xnn_f32_default_params, NULL)
+XNN_UKERNEL_WITH_PARAMS(xnn_arch_riscv_vector, xnn_f32_argmaxpool_ukernel_9x__rvv_u1v, 9, 0, 1, true, float, struct xnn_f32_default_params, NULL)
+
+#endif // XNN_ARCH_RISCV && XNN_ENABLE_RISCV_VECTOR
+
+#ifdef XNN_DEFINED_UKERNEL_WITH_PARAMS
+#undef XNN_DEFINED_UKERNEL_WITH_PARAMS
+#undef XNN_UKERNEL_WITH_PARAMS
+#endif
+
+#ifdef XNN_DEFINED_UKERNEL
+#undef XNN_DEFINED_UKERNEL
+#undef XNN_UKERNEL_UNIPASS
+#endif

src/xnnpack/argmaxpool.h

@@ -14,53 +14,36 @@
 extern "C" {
 #endif
 
-
-#define DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(fn_name) \
+                    
+#define XNN_UKERNEL_UNIPASS(arch_flags,fn_name, primary_tile, incremental_tile, channel_tile, vector_tile, datatype) \
   XNN_INTERNAL void fn_name(                                     \
       size_t output_pixels,                                      \
       size_t kernel_elements,                                    \
       size_t channels,                                           \
-      const float** input,                                       \
+      const datatype** input,                                    \
       size_t input_offset,                                       \
-      float* output,                                             \
+      datatype* output,                                          \
       uint32_t* index,                                           \
       size_t input_increment,                                    \
       size_t output_increment);
+#include "f32-argmaxpool/f32-argmaxpool-unipass.h"
+#undef XNN_UKERNEL_UNIPASS
 
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_4x__neon_c4)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_4x__rvv_u1v)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_4x__scalar_c1)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_4x__sse2_c4)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_4x__wasmsimd_c4)
-
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9x__neon_c4)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9x__rvv_u1v)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9x__scalar_c1)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9x__sse2_c4)
-DECLARE_F32_ARGMAXPOOL_UNIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9x__wasmsimd_c4)
-
-
-#define DECLARE_F32_ARGMAXPOOL_MULTIPASS_UKERNEL_FUNCTION(fn_name) \
+#define XNN_UKERNEL_MULTIPASS(arch_flags,fn_name, primary_tile, incremental_tile, channel_tile, vector_tile, datatype) \
   XNN_INTERNAL void fn_name(                                       \
       size_t output_pixels,                                        \
       size_t kernel_elements,                                      \
       size_t channels,                                             \
-      const float** input,                                         \
+      const datatype** input,                                      \
       size_t input_offset,                                         \
-      float* accumulation_buffer,                                  \
+      datatype* accumulation_buffer,                               \
       uint32_t* index_buffer,                                      \
-      float* output,                                               \
+      datatype* output,                                            \
       uint32_t* index,                                             \
       size_t input_increment,                                      \
       size_t output_increment);
-
-DECLARE_F32_ARGMAXPOOL_MULTIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9p8x__neon_c4)
-DECLARE_F32_ARGMAXPOOL_MULTIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9p8x__rvv_u1v)
-DECLARE_F32_ARGMAXPOOL_MULTIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9p8x__scalar_c1)
-DECLARE_F32_ARGMAXPOOL_MULTIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9p8x__sse2_c4)
-DECLARE_F32_ARGMAXPOOL_MULTIPASS_UKERNEL_FUNCTION(xnn_f32_argmaxpool_ukernel_9p8x__wasmsimd_c4)
-
-
+#include "f32-argmaxpool/f32-argmaxpool-multipass.h"
+#undef XNN_UKERNEL_MULTIPASS
 #ifdef __cplusplus
 }  // extern "C"
 #endif