From fc24dc02cbc5e243b2e728a753ad71ab6d672486 Mon Sep 17 00:00:00 2001 From: Gin890 <79114795+Gin890@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:42:35 +0000 Subject: [PATCH 1/4] Add unit tests for the ScaleInvariantMatrixMatch kernels Add a TestInfra for the four kernel functions (compute_scale and compute_error, plain and weighted). Each test generates a random width x height matrix pair (and weight matrix for the weighted variants), runs the Default path and the ISA path on identical input, and compares the results. The Default path is the source of truth, so each ISA registers the same (width, height, use_weights) combinations against Default. The sizes cover 1x1 up to 1000x1000, both divisible and undivisible by the vector lengths (4, 8, 16), including the short-row cases that exercise unaligned row prefixes and partial vector tails. The SSE4.1 path back-reads up to 3 floats before a row start, so the matrices are allocated with 16 floats of head padding to keep those reads in-bounds. Scales are not computed exactly by all paths (different addition orders), so results are compared with a tolerance: pass if within 4096 ULPs, or within 1e-4 of the magnitude of the accumulated sum (the absolute part covers near-zero results, where ULPs are far smaller than the accumulated rounding of a multi-million-term sum). The tests are registered per-ISA in each x86 Core file (SSE, AVX2, AVX512), guarded by the same PA_AutoDispatch and CPU_CAPABILITY_CURRENT checks the dispatchers use, so each architecture only runs the tests for the ISAs it can execute. --- ...ls_ScaleInvariantMatrixMatch_TestInfra.cpp | 23 ++++++------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp index 463141fd88..7d87bf4191 100644 --- a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp +++ b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp @@ -21,20 +21,11 @@ namespace ScaleInvariantMatrixMatch{ // // Head padding (in floats) before the first row of the matrices. The -// current ISAs do not read before a row start, but the padding keeps the -// first row away from the very start of the allocation in case a future -// ISA does (e.g. an SSE4.1-style back-read of up to VL - 1 floats). +// SSE4.1 implementation back-reads up to 3 floats (VL - 1) before a row +// start to handle unaligned row prefixes, so with this padding those +// reads always stay inside the allocation. // static constexpr size_t HEAD_PAD = 16; -// -// Tail padding after the last row. The SSE4.1 partial-tail load -// (PSHUFB-based `load_f32_no_read_past_end`) always reads a full 16-byte -// word, which extends up to VL - 1 floats past the end of the row even -// though only the first `length` lanes are used. 15 floats (VL - 1 for -// the widest supported ISA) keeps those reads inside the allocation for -// the last row. -// -static constexpr size_t TAIL_PAD = 15; // // A result passes if it is within 4096 ULPs of the Default result, or @@ -103,8 +94,8 @@ UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{ // [0, 1]. TW = T * W is precomputed with the same float rounding // that production uses. // - AlignedVector A(n + HEAD_PAD + TAIL_PAD); - AlignedVector T(n + HEAD_PAD + TAIL_PAD); + AlignedVector A(n + HEAD_PAD); + AlignedVector T(n + HEAD_PAD); AlignedVector rowsA(m_height); AlignedVector rowsT(m_height); AlignedVector W; @@ -112,8 +103,8 @@ UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{ AlignedVector rowsW; AlignedVector rowsTW; if (m_use_weights){ - W = AlignedVector(n + HEAD_PAD + TAIL_PAD); - TW = AlignedVector(n + HEAD_PAD + TAIL_PAD); + W = AlignedVector(n + HEAD_PAD); + TW = AlignedVector(n + HEAD_PAD); rowsW = AlignedVector(m_height); rowsTW = AlignedVector(m_height); } From a16ff872211e6eaa6eb90d730d51ced5325cb217 Mon Sep 17 00:00:00 2001 From: Gin890 <79114795+Gin890@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:38:55 +0000 Subject: [PATCH 2/4] Only register SIMM tests at valid widths and pad the test buffers Each ISA kernel requires a minimum row width (4 for SSE4.1, 8 for AVX2, 16 for AVX512): below it, the partial-row loads span past the row end and can read past the end of the allocation. The test registration now skips sizes below the ISA's minimum width. The test matrices also get 15 floats of tail padding after the last row: the SSE4.1 partial-tail load (PSHUFB-based) reads a full 16-byte word that extends up to VL - 1 floats past the row end, which would otherwise run past the allocation on the last row. --- ...ls_ScaleInvariantMatrixMatch_TestInfra.cpp | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp index 7d87bf4191..463141fd88 100644 --- a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp +++ b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_TestInfra.cpp @@ -21,11 +21,20 @@ namespace ScaleInvariantMatrixMatch{ // // Head padding (in floats) before the first row of the matrices. The -// SSE4.1 implementation back-reads up to 3 floats (VL - 1) before a row -// start to handle unaligned row prefixes, so with this padding those -// reads always stay inside the allocation. +// current ISAs do not read before a row start, but the padding keeps the +// first row away from the very start of the allocation in case a future +// ISA does (e.g. an SSE4.1-style back-read of up to VL - 1 floats). // static constexpr size_t HEAD_PAD = 16; +// +// Tail padding after the last row. The SSE4.1 partial-tail load +// (PSHUFB-based `load_f32_no_read_past_end`) always reads a full 16-byte +// word, which extends up to VL - 1 floats past the end of the row even +// though only the first `length` lanes are used. 15 floats (VL - 1 for +// the widest supported ISA) keeps those reads inside the allocation for +// the last row. +// +static constexpr size_t TAIL_PAD = 15; // // A result passes if it is within 4096 ULPs of the Default result, or @@ -94,8 +103,8 @@ UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{ // [0, 1]. TW = T * W is precomputed with the same float rounding // that production uses. // - AlignedVector A(n + HEAD_PAD); - AlignedVector T(n + HEAD_PAD); + AlignedVector A(n + HEAD_PAD + TAIL_PAD); + AlignedVector T(n + HEAD_PAD + TAIL_PAD); AlignedVector rowsA(m_height); AlignedVector rowsT(m_height); AlignedVector W; @@ -103,8 +112,8 @@ UnitTestResult Test::run(Logger& logger, CancellableScope& scope) const{ AlignedVector rowsW; AlignedVector rowsTW; if (m_use_weights){ - W = AlignedVector(n + HEAD_PAD); - TW = AlignedVector(n + HEAD_PAD); + W = AlignedVector(n + HEAD_PAD + TAIL_PAD); + TW = AlignedVector(n + HEAD_PAD + TAIL_PAD); rowsW = AlignedVector(m_height); rowsTW = AlignedVector(m_height); } From b6b909dcf88c179f8fdc31326de7387313081da9 Mon Sep 17 00:00:00 2001 From: Gin890 <79114795+Gin890@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:47:02 +0000 Subject: [PATCH 3/4] Add the arm64 NEON implementation of the ScaleInvariantMatrixMatch kernels Context_arm64_NEON implements the same SumATA2/SumError routines with float32x4_t vectors (VL = 4). The partial-row loads are lane-insert loads into zeroed vectors, so they never read out of bounds in either direction (no back-reads at the row start, no reads past the row end), unlike the SSE4.1 PSHUFB-based loads. Dispatched for width >= 4 when OK_M1 is enabled; the unit tests (compare against the Default path) are registered for all test sizes with width >= 4. --- .../Kernels_ScaleInvariantMatrixMatch.cpp | 31 +++ ...leInvariantMatrixMatch_Core_arm64_NEON.cpp | 201 ++++++++++++++++++ SerialPrograms/cmake/SourceFiles.cmake | 1 + 3 files changed, 233 insertions(+) create mode 100644 SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp diff --git a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch.cpp b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch.cpp index 0ffedea4fb..576b6ea5ff 100644 --- a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch.cpp +++ b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch.cpp @@ -16,11 +16,13 @@ float compute_scale_Default (size_t width, size_t height, float const* c float compute_scale_min4_x86_SSE (size_t width, size_t height, float const* const* A, float const* const* T); float compute_scale_min8_x86_AVX2 (size_t width, size_t height, float const* const* A, float const* const* T); float compute_scale_min16_x86_AVX512(size_t width, size_t height, float const* const* A, float const* const* T); +float compute_scale_min4_arm64_NEON(size_t width, size_t height, float const* const* A, float const* const* T); float compute_scale_Default (size_t width, size_t height, float const* const* A, float const* const* TW, float const* const* W); float compute_scale_min4_x86_SSE (size_t width, size_t height, float const* const* A, float const* const* TW, float const* const* W); float compute_scale_min8_x86_AVX2 (size_t width, size_t height, float const* const* A, float const* const* TW, float const* const* W); float compute_scale_min16_x86_AVX512(size_t width, size_t height, float const* const* A, float const* const* TW, float const* const* W); +float compute_scale_min4_arm64_NEON(size_t width, size_t height, float const* const* A, float const* const* TW, float const* const* W); float compute_scale( @@ -42,6 +44,11 @@ float compute_scale( if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_08_Nehalem){ return compute_scale_min4_x86_SSE(width, height, A, T); } +#endif +#ifdef PA_AutoDispatch_arm64_20_M1 + if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_M1){ + return compute_scale_min4_arm64_NEON(width, height, A, T); + } #endif return compute_scale_Default(width, height, A, T); } @@ -65,6 +72,11 @@ float compute_scale( if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_08_Nehalem){ return compute_scale_min4_x86_SSE(width, height, A, TW, W); } +#endif +#ifdef PA_AutoDispatch_arm64_20_M1 + if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_M1){ + return compute_scale_min4_arm64_NEON(width, height, A, TW, W); + } #endif return compute_scale_Default(width, height, A, TW, W); } @@ -75,15 +87,18 @@ float compute_error_Default (size_t width, size_t height, float scale, f float compute_error_min4_x86_SSE (size_t width, size_t height, float scale, float const* const* A, float const* const* T); float compute_error_min8_x86_AVX2 (size_t width, size_t height, float scale, float const* const* A, float const* const* T); float compute_error_min16_x86_AVX512(size_t width, size_t height, float scale, float const* const* A, float const* const* T); +float compute_error_min4_arm64_NEON(size_t width, size_t height, float scale, float const* const* A, float const* const* T); float compute_error_Default (size_t width, size_t height, float scale, float const* const* A, float const* const* TW, float const* const* W); float compute_error_min4_x86_SSE (size_t width, size_t height, float scale, float const* const* A, float const* const* TW, float const* const* W); float compute_error_min8_x86_AVX2 (size_t width, size_t height, float scale, float const* const* A, float const* const* TW, float const* const* W); float compute_error_min16_x86_AVX512(size_t width, size_t height, float scale, float const* const* A, float const* const* TW, float const* const* W); +float compute_error_min4_arm64_NEON(size_t width, size_t height, float scale, float const* const* A, float const* const* TW, float const* const* W); void add_tests_x86_SSE(UnitTestDatabase& database); void add_tests_x86_AVX2(UnitTestDatabase& database); void add_tests_x86_AVX512(UnitTestDatabase& database); +void add_tests_arm64_NEON(UnitTestDatabase& database); float compute_error( size_t width, size_t height, @@ -105,6 +120,11 @@ float compute_error( if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_08_Nehalem){ return compute_error_min4_x86_SSE(width, height, scale, A, T); } +#endif +#ifdef PA_AutoDispatch_arm64_20_M1 + if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_M1){ + return compute_error_min4_arm64_NEON(width, height, scale, A, T); + } #endif return compute_error_Default(width, height, scale, A, T); } @@ -129,6 +149,11 @@ float compute_error( if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_08_Nehalem){ return compute_error_min4_x86_SSE(width, height, scale, A, TW, W); } +#endif +#ifdef PA_AutoDispatch_arm64_20_M1 + if (width >= 4 && CPU_CAPABILITY_CURRENT.OK_M1){ + return compute_error_min4_arm64_NEON(width, height, scale, A, TW, W); + } #endif return compute_error_Default(width, height, scale, A, TW, W); } @@ -155,6 +180,12 @@ void add_tests(UnitTestDatabase& database){ add_tests_x86_AVX512(database); } #endif + +#ifdef PA_AutoDispatch_arm64_20_M1 + if (CPU_CAPABILITY_CURRENT.OK_M1){ + add_tests_arm64_NEON(database); + } +#endif } diff --git a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp new file mode 100644 index 0000000000..d098e63cd0 --- /dev/null +++ b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp @@ -0,0 +1,201 @@ +/* Scale Invariant Matrix Match (arm64 NEON) + * + * From: https://github.com/PokemonAutomation/ + * + */ + +#ifdef PA_AutoDispatch_arm64_20_M1 + +#include +#include "Kernels_ScaleInvariantMatrixMatch_Routines.h" +#include "Common/Cpp/TestRunners/UnitTestDatabase.h" +#include "Kernels_ScaleInvariantMatrixMatch_TestInfra.h" + +namespace PokemonAutomation{ +namespace Kernels{ +namespace ScaleInvariantMatrixMatch{ + + +struct Context_arm64_NEON{ + using vtype = float32x4_t; + + static PA_FORCE_INLINE vtype vzero(){ + return vdupq_n_f32(0.0f); + } + static PA_FORCE_INLINE vtype vset(float a){ + return vdupq_n_f32(a); + } + static PA_FORCE_INLINE vtype vadd(vtype a, vtype b){ + return vaddq_f32(a, b); + } + static PA_FORCE_INLINE vtype vmul(vtype a, vtype b){ + return vmulq_f32(a, b); + } + static PA_FORCE_INLINE vtype vpma(vtype a, vtype b, vtype c){ + return vfmaq_f32(c, a, b); + } + static PA_FORCE_INLINE vtype vpms(vtype a, vtype b, vtype c){ + return vfmsq_f32(c, a, b); + } + static PA_FORCE_INLINE float vreduce(vtype a){ + return vaddvq_f32(a); + } + + // + // "ptr" points `length` floats before the row start (the caller shifted + // it back). Load the vector starting at the row start itself and zero + // the last "length" lanes: those row elements are re-counted by the + // next vector (the pointer advances past them), and loading at the row + // start keeps every read in-bounds (the dispatcher only routes rows of + // at least 4 floats to this path, so the 16-byte load stays inside the + // row). No back-read, no read past the end. + // + static PA_FORCE_INLINE vtype load_back(const float* ptr, size_t length){ + const float zero = 0.0f; + vtype out = vld1q_f32(ptr + length); + if (length > 0){ + out = vld1q_lane_f32(&zero, out, 3); + } + if (length > 1){ + out = vld1q_lane_f32(&zero, out, 2); + } + if (length > 2){ + out = vld1q_lane_f32(&zero, out, 1); + } + return out; + } + + // + // "ptr" is valid for exactly "length" floats (1..3): insert them one + // lane at a time into a zeroed vector, so there is no out-of-bounds + // read in either direction. + // + static PA_FORCE_INLINE vtype load_front(const float* ptr, size_t length){ + vtype out = vdupq_n_f32(0.0f); + out = vld1q_lane_f32(ptr, out, 0); + if (length > 1){ + out = vld1q_lane_f32(ptr + 1, out, 1); + } + if (length > 2){ + out = vld1q_lane_f32(ptr + 2, out, 2); + } + return out; + } + + static PA_FORCE_INLINE void load2_partial_back( + size_t length, + vtype& A, const void* ptrA, + vtype& B, const void* ptrB + ){ + A = load_back((const float*)ptrA, length); + B = load_back((const float*)ptrB, length); + } + static PA_FORCE_INLINE void load3_partial_back( + size_t length, + vtype& A, const void* ptrA, + vtype& B, const void* ptrB, + vtype& C, const void* ptrC + ){ + A = load_back((const float*)ptrA, length); + B = load_back((const float*)ptrB, length); + C = load_back((const float*)ptrC, length); + } + static PA_FORCE_INLINE void load2_partial_front( + size_t length, + vtype& A, const void* ptrA, + vtype& B, const void* ptrB + ){ + A = load_front((const float*)ptrA, length); + B = load_front((const float*)ptrB, length); + } + static PA_FORCE_INLINE void load3_partial_front( + size_t length, + vtype& A, const void* ptrA, + vtype& B, const void* ptrB, + vtype& C, const void* ptrC + ){ + A = load_front((const float*)ptrA, length); + B = load_front((const float*)ptrB, length); + C = load_front((const float*)ptrC, length); + } +}; + + + +float compute_scale_min4_arm64_NEON( + size_t width, size_t height, + float const* const* A, + float const* const* T +){ + return compute_scale>(width, height, A, T); +} +float compute_scale_min4_arm64_NEON( + size_t width, size_t height, + float const* const* A, + float const* const* TW, + float const* const* W +){ + return compute_scale>(width, height, A, TW, W); +} +float compute_error_min4_arm64_NEON( + size_t width, size_t height, + float scale, + float const* const* A, + float const* const* T +){ + return compute_error>(width, height, scale, A, T); +} +float compute_error_min4_arm64_NEON( + size_t width, size_t height, + float scale, + float const* const* A, + float const* const* TW, + float const* const* W +){ + return compute_error>(width, height, scale, A, TW, W); +} + + + +float compute_scale_Default(size_t width, size_t height, float const* const* A, float const* const* T); +float compute_scale_Default(size_t width, size_t height, float const* const* A, float const* const* TW, float const* const* W); +float compute_error_Default(size_t width, size_t height, float scale, float const* const* A, float const* const* T); +float compute_error_Default(size_t width, size_t height, float scale, float const* const* A, float const* const* TW, float const* const* W); + +class Test_arm64_NEON : public Test{ +public: + Test_arm64_NEON(size_t width, size_t height, bool use_weights) + : Test( + "arm64_NEON", + width, height, + use_weights, + compute_scale_Default, + compute_scale_min4_arm64_NEON, + compute_scale_Default, + compute_scale_min4_arm64_NEON, + compute_error_Default, + compute_error_min4_arm64_NEON, + compute_error_Default, + compute_error_min4_arm64_NEON + ) + {} +}; +void add_tests_arm64_NEON(UnitTestDatabase& database){ + // The kernel requires a minimum row width (the partial-row loads span up to + // a full vector); sizes below that are not registered for this ISA. + for (size_t i = 0; i < sizeof(kTestSizes) / sizeof(kTestSizes[0]); i++){ + if (kTestSizes[i][0] < 4){ + continue; + } + for (bool use_weights : { false, true }){ + database.add(kTestSizes[i][0], kTestSizes[i][1], use_weights); + } + } +} + + + +} +} +} +#endif diff --git a/SerialPrograms/cmake/SourceFiles.cmake b/SerialPrograms/cmake/SourceFiles.cmake index e662981325..786e0f1cce 100644 --- a/SerialPrograms/cmake/SourceFiles.cmake +++ b/SerialPrograms/cmake/SourceFiles.cmake @@ -1184,6 +1184,7 @@ file(GLOB LIBRARY_SOURCES Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch.cpp Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch.h Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_Default.cpp + Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_x86_AVX2.cpp Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_x86_AVX512.cpp Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_x86_SSE.cpp From 897f0aa372f8970a55abf8c365111c193221d438 Mon Sep 17 00:00:00 2001 From: Alexander Yee Date: Sat, 17 Oct 2026 03:10:12 -0700 Subject: [PATCH 4/4] cleanup --- .../Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp index d098e63cd0..7d3fb31f97 100644 --- a/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp +++ b/SerialPrograms/Source/Kernels/ScaleInvariantMatrixMatch/Kernels_ScaleInvariantMatrixMatch_Core_arm64_NEON.cpp @@ -187,9 +187,8 @@ void add_tests_arm64_NEON(UnitTestDatabase& database){ if (kTestSizes[i][0] < 4){ continue; } - for (bool use_weights : { false, true }){ - database.add(kTestSizes[i][0], kTestSizes[i][1], use_weights); - } + database.add(kTestSizes[i][0], kTestSizes[i][1], false); + database.add(kTestSizes[i][0], kTestSizes[i][1], true); } }