From 9159d7c3d19c0d7e58821fee2aefa91c68abc0a4 Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Wed, 29 Jul 2026 19:44:32 +0300 Subject: [PATCH 1/2] R2 prototype: declarative dispatch table for FP32 IP + FP32 L2 Replaces the #ifdef-nested runtime feature-check cascade in IP_FP32_GetDistFunc/L2_FP32_GetDistFunc with a DispatchTier row struct (dispatch_tier.h) and a pure select_tier_index helper that only does predicate/min_dim comparisons - it never calls a row's chooser, so it's safe to exercise with fabricated Features values regardless of what the host CPU supports, unlike the full GetDistFunc/Choose_* path which calls into an ISA-specific translation unit. The actual per-combo tables (IP_FP32_DispatchTable, L2_FP32_DispatchTable) live in new headers, not as function-locals, because inline constexpr variables need their initializer visible in every TU that reads them - test_spaces.cpp and the real GetDistFunc consume the exact same table object, not a hand-duplicated copy that could drift. Exposes spaces::FeaturesType (was a function-local alias inside getCpuOptimizationFeatures) since DispatchTier's predicate signature needs a name for it. Verified: standalone compile of IP_space.cpp/L2_space.cpp/test_spaces.cpp; a mock-table check of select_tier_index's priority ordering and min_dim gating in isolation; new gtest cases exercising the real production tables via fabricated features (no-feature fallback, overlapping-feature priority, missing-top-feature fallthrough, min_dim-1/min_dim/min_dim+1 boundaries). 13 more GetDistFunc functions remain in IP_space.cpp, 8 in L2_space.cpp. --- src/VecSim/spaces/IP_dispatch_tables.h | 65 +++++++++++++++ src/VecSim/spaces/IP_space.cpp | 58 ++----------- src/VecSim/spaces/L2_dispatch_tables.h | 56 +++++++++++++ src/VecSim/spaces/L2_space.cpp | 57 ++----------- src/VecSim/spaces/dispatch_tier.h | 64 ++++++++++++++ src/VecSim/spaces/spaces.h | 9 +- tests/unit/test_spaces.cpp | 111 +++++++++++++++++++++++++ 7 files changed, 320 insertions(+), 100 deletions(-) create mode 100644 src/VecSim/spaces/IP_dispatch_tables.h create mode 100644 src/VecSim/spaces/L2_dispatch_tables.h create mode 100644 src/VecSim/spaces/dispatch_tier.h diff --git a/src/VecSim/spaces/IP_dispatch_tables.h b/src/VecSim/spaces/IP_dispatch_tables.h new file mode 100644 index 000000000..7c9c3a102 --- /dev/null +++ b/src/VecSim/spaces/IP_dispatch_tables.h @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ +#pragma once + +#include "VecSim/spaces/dispatch_tier.h" +#include "VecSim/spaces/functions/SVE2.h" +#include "VecSim/spaces/functions/SVE.h" +#include "VecSim/spaces/functions/NEON.h" +#include "VecSim/spaces/functions/AVX512F.h" +#include "VecSim/spaces/functions/AVX.h" +#include "VecSim/spaces/functions/SSE.h" + +/* + * `inline constexpr` dispatch tables for IP_space.cpp's GetDistFunc-per-combo functions, one per + * migrated (dtype, metric) combo. Defined here (not in IP_space.cpp) because both IP_space.cpp + * and test_spaces.cpp need to see the same table object with its full initializer visible - + * `inline constexpr` variables need their definition present in every translation unit that uses + * them, so a header is the only place this can live. Row order matches the priority order of the + * `#ifdef` cascade this replaces exactly (best tier first); each row is still individually + * `#ifdef OPT_`-guarded, so the table only ever contains rows for tiers this toolchain can + * actually emit - the compile-time file-gating this refactor deliberately keeps. + */ +namespace spaces { + +// Deliberately size-deduced (std::array CTAD), not a fixed <..., 6> - the actual row count varies +// with which OPT_* tiers this toolchain/target compiles in (e.g. this table has only 3 rows on +// an x86 build with no AVX512F, and a different 3 on ARM), so a fixed size would either fail to +// compile or (worse) silently zero-initialize missing rows into null-predicate garbage. +inline constexpr auto IP_FP32_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_FP32_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_FP32_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 0, 0, + Choose_FP32_IP_implementation_NEON}, +#endif +#ifdef OPT_AVX512F + // Optimizations assume at least 8 floats (see the residual handling in the kernels); below + // that, the scalar implementation is at least as fast anyway - hence min_dim=8 on every x86 + // row here, matching the shared `if (dim < 8) return ret_dist_func;` gate it replaces. + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 8, 16, + Choose_FP32_IP_implementation_AVX512F}, +#endif +#ifdef OPT_AVX + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx; }, 8, 8, + Choose_FP32_IP_implementation_AVX}, +#endif +#ifdef OPT_SSE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse; }, 8, 4, + Choose_FP32_IP_implementation_SSE}, +#endif +}; + +} // namespace spaces diff --git a/src/VecSim/spaces/IP_space.cpp b/src/VecSim/spaces/IP_space.cpp index 6706d8f31..f6dc95f53 100644 --- a/src/VecSim/spaces/IP_space.cpp +++ b/src/VecSim/spaces/IP_space.cpp @@ -9,6 +9,7 @@ #include "VecSim/spaces/space_includes.h" #include "VecSim/spaces/IP_space.h" #include "VecSim/spaces/IP/IP.h" +#include "VecSim/spaces/IP_dispatch_tables.h" #include "VecSim/types/bfloat16.h" #include "VecSim/types/float16.h" #include "VecSim/spaces/functions/AVX512F.h" @@ -426,57 +427,16 @@ dist_func_t IP_FP32_GetDistFunc(size_t dim, unsigned char *alignment, con alignment = &dummy_alignment; } - dist_func_t ret_dist_func = FP32_InnerProduct; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); -#ifdef CPU_FEATURES_ARCH_AARCH64 - -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_FP32_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_FP32_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd) { - return Choose_FP32_IP_implementation_NEON(dim); - } -#endif - -#endif - -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 8 floats (see the residual handling in the kernels). - // Below that, the scalar implementation is at least as fast anyway. - if (dim < 8) { - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_FP32_DispatchTable); + if (idx == IP_FP32_DispatchTable.size()) { + return FP32_InnerProduct; } -#ifdef OPT_AVX512F - if (features.avx512f) { - if (dim % 16 == 0) // no point in aligning if we have an offsetting residual - *alignment = 16 * sizeof(float); // handles 16 floats - return Choose_FP32_IP_implementation_AVX512F(dim); + const auto &tier = IP_FP32_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(float); } -#endif -#ifdef OPT_AVX - if (features.avx) { - if (dim % 8 == 0) // no point in aligning if we have an offsetting residual - *alignment = 8 * sizeof(float); // handles 8 floats - return Choose_FP32_IP_implementation_AVX(dim); - } -#endif -#ifdef OPT_SSE - if (features.sse) { - if (dim % 4 == 0) // no point in aligning if we have an offsetting residual - *alignment = 4 * sizeof(float); // handles 4 floats - return Choose_FP32_IP_implementation_SSE(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } dist_func_t IP_FP64_GetDistFunc(size_t dim, unsigned char *alignment, diff --git a/src/VecSim/spaces/L2_dispatch_tables.h b/src/VecSim/spaces/L2_dispatch_tables.h new file mode 100644 index 000000000..c08489ab8 --- /dev/null +++ b/src/VecSim/spaces/L2_dispatch_tables.h @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ +#pragma once + +#include "VecSim/spaces/dispatch_tier.h" +#include "VecSim/spaces/functions/SVE2.h" +#include "VecSim/spaces/functions/SVE.h" +#include "VecSim/spaces/functions/NEON.h" +#include "VecSim/spaces/functions/AVX512F.h" +#include "VecSim/spaces/functions/AVX.h" +#include "VecSim/spaces/functions/SSE.h" + +/* + * `inline constexpr` dispatch tables for L2_space.cpp's GetDistFunc-per-combo functions - see + * IP_dispatch_tables.h for the full rationale (same reasons apply here, mirrored for L2). + */ +namespace spaces { + +// Deliberately size-deduced (std::array CTAD), not a fixed size - see IP_dispatch_tables.h. +inline constexpr auto L2_FP32_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_FP32_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_FP32_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 0, 0, + Choose_FP32_L2_implementation_NEON}, +#endif +#ifdef OPT_AVX512F + // Optimizations assume at least 8 floats (see the residual handling in the kernels); below + // that, the scalar implementation is at least as fast anyway - hence min_dim=8 on every x86 + // row here, matching the shared `if (dim < 8) return ret_dist_func;` gate it replaces. + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 8, 16, + Choose_FP32_L2_implementation_AVX512F}, +#endif +#ifdef OPT_AVX + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx; }, 8, 8, + Choose_FP32_L2_implementation_AVX}, +#endif +#ifdef OPT_SSE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse; }, 8, 4, + Choose_FP32_L2_implementation_SSE}, +#endif +}; + +} // namespace spaces diff --git a/src/VecSim/spaces/L2_space.cpp b/src/VecSim/spaces/L2_space.cpp index 07e638cba..d048a0c99 100644 --- a/src/VecSim/spaces/L2_space.cpp +++ b/src/VecSim/spaces/L2_space.cpp @@ -9,6 +9,7 @@ #include "VecSim/spaces/space_includes.h" #include "VecSim/spaces/L2_space.h" #include "VecSim/spaces/L2/L2.h" +#include "VecSim/spaces/L2_dispatch_tables.h" #include "VecSim/types/bfloat16.h" #include "VecSim/types/float16.h" #include "VecSim/spaces/functions/F16C.h" @@ -188,56 +189,16 @@ dist_func_t L2_FP32_GetDistFunc(size_t dim, unsigned char *alignment, con alignment = &dummy_alignment; } - dist_func_t ret_dist_func = FP32_L2Sqr; - - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_FP32_L2_implementation_SVE2(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_FP32_DispatchTable); + if (idx == L2_FP32_DispatchTable.size()) { + return FP32_L2Sqr; } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_FP32_L2_implementation_SVE(dim); + const auto &tier = L2_FP32_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(float); } -#endif -#ifdef OPT_NEON - if (features.asimd) { - return Choose_FP32_L2_implementation_NEON(dim); - } -#endif -#endif - -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 8 floats (see the residual handling in the kernels). - // Below that, the scalar implementation is at least as fast anyway. - if (dim < 8) { - return ret_dist_func; - } -#ifdef OPT_AVX512F - if (features.avx512f) { - if (dim % 16 == 0) // no point in aligning if we have an offsetting residual - *alignment = 16 * sizeof(float); // handles 16 floats - return Choose_FP32_L2_implementation_AVX512F(dim); - } -#endif -#ifdef OPT_AVX - if (features.avx) { - if (dim % 8 == 0) // no point in aligning if we have an offsetting residual - *alignment = 8 * sizeof(float); // handles 8 floats - return Choose_FP32_L2_implementation_AVX(dim); - } -#endif -#ifdef OPT_SSE - if (features.sse) { - if (dim % 4 == 0) // no point in aligning if we have an offsetting residual - *alignment = 4 * sizeof(float); // handles 4 floats - return Choose_FP32_L2_implementation_SSE(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } dist_func_t L2_FP64_GetDistFunc(size_t dim, unsigned char *alignment, diff --git a/src/VecSim/spaces/dispatch_tier.h b/src/VecSim/spaces/dispatch_tier.h new file mode 100644 index 000000000..c4efc214f --- /dev/null +++ b/src/VecSim/spaces/dispatch_tier.h @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ +#pragma once + +#include "VecSim/spaces/spaces.h" +#include +#include + +/* + * One row of a per-(dtype,metric) dispatch table, replacing a single `#ifdef OPT_ if + * (features.) { ... return Choose_*(dim); }` block from the old cascade with one struct + * literal. Rows are listed in priority order (best tier first) inside a + * `inline constexpr std::array, N>` defined in a dispatch-tables header + * (e.g. IP_dispatch_tables.h) so both the real GetDistFunc and test_spaces.cpp read the exact + * same table. + */ +namespace spaces { + +template +struct DispatchTier { + // Pure CPU-feature check only - no dimension logic here, see min_dim below. Must be a plain + // function pointer (a captureless lambda decays to one), never std::function - this sits on + // a cold path but there's no reason to pay for type erasure/possible allocation anyway. + bool (*predicate)(const FeaturesType &) noexcept; + // Dimension floor below which this tier is skipped even if predicate matches. 0 = no floor. + size_t min_dim; + // Alignment hint, in elements of the tier's storage type (not bytes - the caller multiplies + // by sizeof(StorageElemType), which this struct doesn't know). 0 = this tier never sets an + // alignment hint (the two documented cosine skip-alignment special cases use this). + size_t alignment_chunk_elems; + // The existing Choose_* function for this tier - unchanged by this refactor. + dist_func_t (*chooser)(size_t dim); +}; + +/* + * Returns the index of the first row in `rows` whose predicate matches `features` and whose + * min_dim is satisfied by `dim`, or `rows.size()` if no row matches (caller falls back to the + * naive/scalar implementation, exactly as the old cascade's final `return ret_dist_func;` did). + * + * Pure comparisons only - this never calls any row's `chooser`, and therefore never touches any + * arch-specific translation unit. That makes it safe to call with an arbitrary or fabricated + * `FeaturesType` value regardless of what the host CPU actually supports - unlike the full + * GetDistFunc path, which calls into a TU compiled for a specific ISA target and so must only + * ever be exercised with real, host-detected features (see SPACES-REFACTOR-PLAN.md's + * execution-safety rule). + */ +template +size_t select_tier_index(const FeaturesType &features, size_t dim, + const std::array, N> &rows) { + for (size_t i = 0; i < N; i++) { + if (rows[i].predicate(features) && dim >= rows[i].min_dim) { + return i; + } + } + return N; +} + +} // namespace spaces diff --git a/src/VecSim/spaces/spaces.h b/src/VecSim/spaces/spaces.h index 11b0f9801..abdb24960 100644 --- a/src/VecSim/spaces/spaces.h +++ b/src/VecSim/spaces/spaces.h @@ -52,13 +52,16 @@ static int inline is_little_endian() { return *(char *)&x; } -static inline auto getCpuOptimizationFeatures(const void *arch_opt = nullptr) { +#if defined(CPU_FEATURES_ARCH_AARCH64) +using FeaturesType = cpu_features::Aarch64Features; +#else +using FeaturesType = cpu_features::X86Features; // Fallback +#endif +static inline auto getCpuOptimizationFeatures(const void *arch_opt = nullptr) { #if defined(CPU_FEATURES_ARCH_AARCH64) - using FeaturesType = cpu_features::Aarch64Features; constexpr auto getFeatures = cpu_features::GetAarch64Info; #else - using FeaturesType = cpu_features::X86Features; // Fallback constexpr auto getFeatures = cpu_features::GetX86Info; #endif return arch_opt ? *static_cast(arch_opt) : getFeatures().features; diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 8e53c83d1..2056e393b 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -22,6 +22,8 @@ #include "VecSim/types/bfloat16.h" #include "VecSim/spaces/IP_space.h" #include "VecSim/spaces/L2_space.h" +#include "VecSim/spaces/IP_dispatch_tables.h" +#include "VecSim/spaces/L2_dispatch_tables.h" #include "VecSim/types/float16.h" #include "VecSim/types/sq8.h" #include "VecSim/spaces/functions/AVX512F.h" @@ -877,6 +879,115 @@ TEST_P(FP32SpacesOptimizationTest, FP32InnerProductTest) { INSTANTIATE_TEST_SUITE_P(FP32OptFuncs, FP32SpacesOptimizationTest, testing::Range(8UL, 32 * 2UL + 1)); +// select_tier_index is pure comparison logic with no dependency on any arch-specific translation +// unit (it never calls a row's `chooser`), so - unlike the Choose_*/GetDistFunc tests above, which +// must only ever use real, host-detected features - these tests are free to fabricate any +// FeaturesType bit pattern, including combinations the host running this test doesn't actually +// have. They exercise the *real* production IP_FP32_DispatchTable/L2_FP32_DispatchTable (not a +// hand-duplicated copy), so a change to the real table's row order/predicates/min_dim is caught +// here even though it can never crash: no chooser is ever invoked. +TEST(SpacesDispatchTierTest, FP32IPTableNoFeaturesFallsBackToScalar) { + spaces::FeaturesType none{}; + EXPECT_EQ(spaces::select_tier_index(none, 1000, spaces::IP_FP32_DispatchTable), + spaces::IP_FP32_DispatchTable.size()); +} + +TEST(SpacesDispatchTierTest, FP32L2TableNoFeaturesFallsBackToScalar) { + spaces::FeaturesType none{}; + EXPECT_EQ(spaces::select_tier_index(none, 1000, spaces::L2_FP32_DispatchTable), + spaces::L2_FP32_DispatchTable.size()); +} + +#ifdef CPU_FEATURES_ARCH_X86_64 +TEST(SpacesDispatchTierTest, FP32IPTableOverlappingFeaturePriority) { + spaces::FeaturesType all{}; + all.avx512f = 1; + all.avx = 1; + all.sse = 1; + // Best tier first: avx512f must win over avx and sse even though all three predicates match. + size_t idx = spaces::select_tier_index(all, 1000, spaces::IP_FP32_DispatchTable); + ASSERT_NE(idx, spaces::IP_FP32_DispatchTable.size()); + EXPECT_EQ(spaces::IP_FP32_DispatchTable[idx].chooser, + spaces::Choose_FP32_IP_implementation_AVX512F); +} + +TEST(SpacesDispatchTierTest, FP32IPTableMissingTopFeatureFallsThrough) { + spaces::FeaturesType avx_and_sse{}; + avx_and_sse.avx = 1; + avx_and_sse.sse = 1; + size_t idx = spaces::select_tier_index(avx_and_sse, 1000, spaces::IP_FP32_DispatchTable); + ASSERT_NE(idx, spaces::IP_FP32_DispatchTable.size()); + EXPECT_EQ(spaces::IP_FP32_DispatchTable[idx].chooser, + spaces::Choose_FP32_IP_implementation_AVX); +} + +TEST(SpacesDispatchTierTest, FP32IPTableOnlyWeakestFeature) { + spaces::FeaturesType sse_only{}; + sse_only.sse = 1; + size_t idx = spaces::select_tier_index(sse_only, 1000, spaces::IP_FP32_DispatchTable); + ASSERT_NE(idx, spaces::IP_FP32_DispatchTable.size()); + EXPECT_EQ(spaces::IP_FP32_DispatchTable[idx].chooser, + spaces::Choose_FP32_IP_implementation_SSE); +} + +// Every x86 row in this table shares min_dim=8 (the old shared "if (dim < 8) return scalar;" +// gate, absorbed per-row - see IP_dispatch_tables.h). min_dim-1/min_dim/min_dim+1 boundary. +TEST(SpacesDispatchTierTest, FP32IPTableMinDimBoundary) { + spaces::FeaturesType avx512{}; + avx512.avx512f = 1; + EXPECT_EQ(spaces::select_tier_index(avx512, 7, spaces::IP_FP32_DispatchTable), + spaces::IP_FP32_DispatchTable.size()) + << "dim=7 is below min_dim=8, must fall back to scalar even with avx512f set"; + size_t idx8 = spaces::select_tier_index(avx512, 8, spaces::IP_FP32_DispatchTable); + ASSERT_NE(idx8, spaces::IP_FP32_DispatchTable.size()); + EXPECT_EQ(spaces::IP_FP32_DispatchTable[idx8].chooser, + spaces::Choose_FP32_IP_implementation_AVX512F); + size_t idx9 = spaces::select_tier_index(avx512, 9, spaces::IP_FP32_DispatchTable); + ASSERT_NE(idx9, spaces::IP_FP32_DispatchTable.size()); + EXPECT_EQ(spaces::IP_FP32_DispatchTable[idx9].chooser, + spaces::Choose_FP32_IP_implementation_AVX512F); +} + +TEST(SpacesDispatchTierTest, FP32L2TableOverlappingFeaturePriority) { + spaces::FeaturesType all{}; + all.avx512f = 1; + all.avx = 1; + all.sse = 1; + size_t idx = spaces::select_tier_index(all, 1000, spaces::L2_FP32_DispatchTable); + ASSERT_NE(idx, spaces::L2_FP32_DispatchTable.size()); + EXPECT_EQ(spaces::L2_FP32_DispatchTable[idx].chooser, + spaces::Choose_FP32_L2_implementation_AVX512F); +} +#endif // CPU_FEATURES_ARCH_X86_64 + +#ifdef CPU_FEATURES_ARCH_AARCH64 +TEST(SpacesDispatchTierTest, FP32IPTableOverlappingFeaturePriority) { + spaces::FeaturesType all{}; + all.sve2 = 1; + all.sve = 1; + all.asimd = 1; + // Best tier first: sve2 must win over sve and neon even though all three predicates match. + size_t idx = spaces::select_tier_index(all, 1000, spaces::IP_FP32_DispatchTable); + ASSERT_NE(idx, spaces::IP_FP32_DispatchTable.size()); + EXPECT_EQ(spaces::IP_FP32_DispatchTable[idx].chooser, + spaces::Choose_FP32_IP_implementation_SVE2); +} + +// Every ARM row in this table has min_dim=0 (no shared dimension gate on this path in the +// original cascade), so there is no min_dim-1 underflow case to test - dims 0 and 1 instead. +TEST(SpacesDispatchTierTest, FP32IPTableZeroMinDimAcceptsSmallestDims) { + spaces::FeaturesType neon_only{}; + neon_only.asimd = 1; + for (size_t dim : {size_t{0}, size_t{1}}) { + size_t idx = spaces::select_tier_index(neon_only, dim, spaces::IP_FP32_DispatchTable); + ASSERT_NE(idx, spaces::IP_FP32_DispatchTable.size()) << "dim=" << dim; + EXPECT_EQ(spaces::IP_FP32_DispatchTable[idx].chooser, + spaces::Choose_FP32_IP_implementation_NEON) + << "dim=" << dim; + } +} +#endif // CPU_FEATURES_ARCH_AARCH64 + class FP64SpacesOptimizationTest : public testing::TestWithParam {}; TEST_P(FP64SpacesOptimizationTest, FP64L2SqrTest) { From 06f58f41fb7fe3320eb3194ca0e923f5d8dc473b Mon Sep 17 00:00:00 2001 From: Dor Forer Date: Wed, 29 Jul 2026 21:51:56 +0300 Subject: [PATCH 2/2] Finish R2: declarative dispatch tables for all remaining GetDistFunc combos Converts the remaining 21 GetDistFunc functions (13 in IP_space.cpp, 8 in L2_space.cpp) from #ifdef-nested runtime feature-check cascades to the DispatchTier/select_tier_index pattern established by the FP32 prototype. Per-row predicate/min_dim/alignment_chunk_elems were transcribed from a fresh read of each original cascade, not from memory, to avoid transcription drift on the two special cases in this codebase: - IP's Cosine_INT8/Cosine_UINT8 AVX-512 tiers deliberately never set an alignment hint (alignment_chunk_elems=0) - the extra norm float shifts effective alignment in a way the original code skips computing, to avoid complexity. There is no L2 mirror: L2_INT8/L2_UINT8 set alignment normally. - IP_BF16 has an AVX512BF16_VL tier that L2_BF16 never had in the original cascade (L2 only ever dispatched BF16 to AVX512BW_VBMI2 on that rung). BF16's big/little-endian split is not a tier-selection concern and stays a special-cased branch before the table is consulted, in both files. Verified: standalone compile of both _space.cpp files and test_spaces.cpp; zero remaining `if (features.` cascades (grep-confirmed); new gtest cases asserting both special cases above survive in the actual production tables via the pure, ISA-independent select_tier_index (safe to fabricate features against - never calls a chooser). Existing exhaustive real-hardware tests (BF16/INT8/UINT8/SQ8*OptFuncs etc.) are unchanged and remain the regression oracle for full GetDistFunc behavior once built. --- src/VecSim/spaces/IP_dispatch_tables.h | 425 ++++++++++++++ src/VecSim/spaces/IP_space.cpp | 756 ++++--------------------- src/VecSim/spaces/L2_dispatch_tables.h | 257 ++++++++- src/VecSim/spaces/L2_space.cpp | 432 +++----------- tests/unit/test_spaces.cpp | 71 +++ 5 files changed, 942 insertions(+), 999 deletions(-) diff --git a/src/VecSim/spaces/IP_dispatch_tables.h b/src/VecSim/spaces/IP_dispatch_tables.h index 7c9c3a102..4ebb84f1a 100644 --- a/src/VecSim/spaces/IP_dispatch_tables.h +++ b/src/VecSim/spaces/IP_dispatch_tables.h @@ -11,10 +11,26 @@ #include "VecSim/spaces/dispatch_tier.h" #include "VecSim/spaces/functions/SVE2.h" #include "VecSim/spaces/functions/SVE.h" +#include "VecSim/spaces/functions/SVE_BF16.h" #include "VecSim/spaces/functions/NEON.h" +#include "VecSim/spaces/functions/NEON_DOTPROD.h" +#include "VecSim/spaces/functions/NEON_HP.h" +#include "VecSim/spaces/functions/NEON_BF16.h" #include "VecSim/spaces/functions/AVX512F.h" +#include "VecSim/spaces/functions/AVX512BW_VBMI2.h" +#include "VecSim/spaces/functions/AVX512BF16_VL.h" +#include "VecSim/spaces/functions/AVX512FP16_VL.h" +#include "VecSim/spaces/functions/AVX512F_BW_VL_VNNI.h" #include "VecSim/spaces/functions/AVX.h" +#include "VecSim/spaces/functions/AVX2.h" +#include "VecSim/spaces/functions/AVX2_F16C.h" +#include "VecSim/spaces/functions/AVX2_FMA.h" +#include "VecSim/spaces/functions/AVX2_FMA_F16C.h" +#include "VecSim/spaces/functions/F16C.h" #include "VecSim/spaces/functions/SSE.h" +#include "VecSim/spaces/functions/SSE3.h" +#include "VecSim/spaces/functions/SSE4.h" +#include "VecSim/spaces/functions/SSE4_F16C.h" /* * `inline constexpr` dispatch tables for IP_space.cpp's GetDistFunc-per-combo functions, one per @@ -25,6 +41,10 @@ * `#ifdef` cascade this replaces exactly (best tier first); each row is still individually * `#ifdef OPT_`-guarded, so the table only ever contains rows for tiers this toolchain can * actually emit - the compile-time file-gating this refactor deliberately keeps. + * + * Every table here is deliberately size-deduced (`std::array` CTAD), never a fixed `<..., N>` - + * see IP_FP32_DispatchTable's comment for why a fixed size would be actively dangerous, not just + * inconvenient. */ namespace spaces { @@ -62,4 +82,409 @@ inline constexpr auto IP_FP32_DispatchTable = std::array{ #endif }; +// Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. +inline constexpr auto IP_SQ8_FP32_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_SQ8_FP32_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_SQ8_FP32_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 0, 0, + Choose_SQ8_FP32_IP_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vnni); + }, + 8, 16, Choose_SQ8_FP32_IP_implementation_AVX512F_BW_VL_VNNI}, +#endif +#ifdef OPT_AVX2_FMA + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.fma3); }, 8, + 8, Choose_SQ8_FP32_IP_implementation_AVX2_FMA}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx2; }, 8, 8, + Choose_SQ8_FP32_IP_implementation_AVX2}, +#endif +#ifdef OPT_SSE4 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse4_1; }, 8, 4, + Choose_SQ8_FP32_IP_implementation_SSE4}, +#endif +}; + +inline constexpr auto Cosine_SQ8_FP32_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_SQ8_FP32_Cosine_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_SQ8_FP32_Cosine_implementation_SVE}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 0, 0, + Choose_SQ8_FP32_Cosine_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vnni); + }, + 8, 16, Choose_SQ8_FP32_Cosine_implementation_AVX512F_BW_VL_VNNI}, +#endif +#ifdef OPT_AVX2_FMA + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.fma3); }, 8, + 8, Choose_SQ8_FP32_Cosine_implementation_AVX2_FMA}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx2; }, 8, 8, + Choose_SQ8_FP32_Cosine_implementation_AVX2}, +#endif +#ifdef OPT_SSE4 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse4_1; }, 8, 4, + Choose_SQ8_FP32_Cosine_implementation_SSE4}, +#endif +}; + +// SQ8<->FP16 tiers all need F16C (vcvtph2ps) except AVX-512 (cvtph_ps is part of AVX-512F +// itself). min_dim=16 on every row - both the x86 and the ARM `#ifdef` blocks in the original +// cascade shared one `if (dim < 16) return ret_dist_func;` gate over all their rows. +inline constexpr auto IP_SQ8_FP16_DispatchTable = std::array{ +#ifdef OPT_AVX512F + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 16, 16, + Choose_SQ8_FP16_IP_implementation_AVX512F}, +#endif +#ifdef OPT_F16C +#ifdef OPT_AVX2_FMA + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.fma3 && f.f16c); }, 16, 8, + Choose_SQ8_FP16_IP_implementation_AVX2_FMA}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.f16c); }, 16, + 8, Choose_SQ8_FP16_IP_implementation_AVX2}, +#endif +#ifdef OPT_SSE4 + // F16C is VEX-encoded - require AVX as well, matching the existing F16C/FP16 dispatcher. + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.sse4_1 && f.f16c && f.avx); }, 16, 4, + Choose_SQ8_FP16_IP_implementation_SSE4}, +#endif +#endif // OPT_F16C +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 16, 0, + Choose_SQ8_FP16_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 16, 0, + Choose_SQ8_FP16_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON_HP + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdfhm; }, 16, 0, + Choose_SQ8_FP16_IP_implementation_NEON_FHM}, + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdhp; }, 16, 0, + Choose_SQ8_FP16_IP_implementation_NEON_HP}, +#endif +}; + +inline constexpr auto Cosine_SQ8_FP16_DispatchTable = std::array{ +#ifdef OPT_AVX512F + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 16, 16, + Choose_SQ8_FP16_Cosine_implementation_AVX512F}, +#endif +#ifdef OPT_F16C +#ifdef OPT_AVX2_FMA + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.fma3 && f.f16c); }, 16, 8, + Choose_SQ8_FP16_Cosine_implementation_AVX2_FMA}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.f16c); }, 16, + 8, Choose_SQ8_FP16_Cosine_implementation_AVX2}, +#endif +#ifdef OPT_SSE4 + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.sse4_1 && f.f16c && f.avx); }, 16, 4, + Choose_SQ8_FP16_Cosine_implementation_SSE4}, +#endif +#endif // OPT_F16C +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 16, 0, + Choose_SQ8_FP16_Cosine_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 16, 0, + Choose_SQ8_FP16_Cosine_implementation_SVE}, +#endif +#ifdef OPT_NEON_HP + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdfhm; }, 16, 0, + Choose_SQ8_FP16_Cosine_implementation_NEON_FHM}, + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdhp; }, 16, 0, + Choose_SQ8_FP16_Cosine_implementation_NEON_HP}, +#endif +}; + +// Both operands are SQ8 with a precomputed sum; the ARM rows carry their own per-row dim>=16 +// floor (no shared gate existed for them in the original cascade), and the single x86 tier uses +// 64-element chunks (residual handling is in 32-byte sub-chunks) with its own dim>=64 floor. +inline constexpr auto IP_SQ8_SQ8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_SQ8_SQ8_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_SQ8_SQ8_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_SQ8_SQ8_IP_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_SQ8_SQ8_IP_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vnni); + }, + 64, 32, Choose_SQ8_SQ8_IP_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + +inline constexpr auto Cosine_SQ8_SQ8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_SQ8_SQ8_Cosine_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_SQ8_SQ8_Cosine_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_SQ8_SQ8_Cosine_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_SQ8_SQ8_Cosine_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vnni); + }, + 64, 32, Choose_SQ8_SQ8_Cosine_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + +inline constexpr auto IP_FP64_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_FP64_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_FP64_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 0, 0, + Choose_FP64_IP_implementation_NEON}, +#endif +#ifdef OPT_AVX512F + // Optimizations assume at least 4 doubles; below that, scalar is at least as fast anyway. + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 4, 8, + Choose_FP64_IP_implementation_AVX512F}, +#endif +#ifdef OPT_AVX + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx; }, 4, 4, + Choose_FP64_IP_implementation_AVX}, +#endif +#ifdef OPT_SSE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse; }, 4, 2, + Choose_FP64_IP_implementation_SSE}, +#endif +}; + +// The BF16 big/little-endian split is not a tier-selection concern at all (it's not a CPU-feature +// or dimension decision) and stays as a special-cased branch in IP_BF16_GetDistFunc itself, +// before this table is ever consulted - see IP_space.cpp. +inline constexpr auto IP_BF16_DispatchTable = std::array{ +#ifdef OPT_SVE_BF16 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.svebf16; }, 0, 0, + Choose_BF16_IP_implementation_SVE_BF16}, +#endif +#ifdef OPT_NEON_BF16 + // Optimization assumes at least 8 BF16s (a full chunk). + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.bf16; }, 8, 0, + Choose_BF16_IP_implementation_NEON_BF16}, +#endif +#ifdef OPT_AVX512_BF16_VL + // Optimizations assume at least 32 bfloats; below that, scalar is at least as fast anyway. + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx512_bf16 && f.avx512vl); }, 32, 32, + Choose_BF16_IP_implementation_AVX512BF16_VL}, +#endif +#ifdef OPT_AVX512_BW_VBMI2 + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx512bw && f.avx512vbmi2); }, 32, 32, + Choose_BF16_IP_implementation_AVX512BW_VBMI2}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx2; }, 32, 16, + Choose_BF16_IP_implementation_AVX2}, +#endif +#ifdef OPT_SSE3 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse3; }, 32, 8, + Choose_BF16_IP_implementation_SSE3}, +#endif +}; + +// Each x86/ARM tier here carries its own per-row dimension floor, implied by its residual +// handling (see the original cascade's comment in IP_space.cpp for the exact rationale per tier) +// - there was no single shared gate to absorb, unlike IP_FP32/IP_SQ8_FP16/etc. +inline constexpr auto IP_FP16_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_FP16_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_FP16_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON_HP + // Optimization assumes at least 8 16FPs (a full chunk). + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdhp; }, 8, 0, + Choose_FP16_IP_implementation_NEON_HP}, +#endif +#ifdef OPT_AVX512_FP16_VL + // The AVX512FP16_VL kernel loads full 512-bit blocks (32 elements). + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx512_fp16 && f.avx512vl); }, 32, 32, + Choose_FP16_IP_implementation_AVX512FP16_VL}, +#endif +#ifdef OPT_AVX512F + // The AVX512F kernel loads full 256-bit blocks (16 elements). + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 16, 32, + Choose_FP16_IP_implementation_AVX512F}, +#endif +#ifdef OPT_F16C + // The F16C kernel loads full 128-bit blocks (8 elements). + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.f16c && f.fma3 && f.avx); }, 8, 16, + Choose_FP16_IP_implementation_F16C}, +#endif +}; + +inline constexpr auto IP_INT8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_INT8_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_INT8_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_INT8_IP_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_INT8_IP_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + // Optimizations assume at least 32 int8; below that, scalar is at least as fast anyway. + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vl && f.avx512vnni); + }, + 32, 32, Choose_INT8_IP_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + +// Cosine's AVX-512 tier deliberately never sets an alignment hint (alignment_chunk_elems=0): for +// int8 vectors with cosine distance, the extra float for the norm shifts the effective alignment +// to `(dim + sizeof(float)) % 32`, and vectors satisfying THAT have a residual, causing offset +// loads during calculation. The original cascade skips computing this to avoid the complexity, +// assuming the performance impact is negligible - carried over unchanged. There is no L2 mirror +// of this special case: L2_INT8_GetDistFunc's AVX-512 tier sets alignment normally. +inline constexpr auto Cosine_INT8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_INT8_Cosine_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_INT8_Cosine_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_INT8_Cosine_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_INT8_Cosine_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vl && f.avx512vnni); + }, + 32, 0, Choose_INT8_Cosine_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + +inline constexpr auto IP_UINT8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_UINT8_IP_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_UINT8_IP_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_UINT8_IP_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_UINT8_IP_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vl && f.avx512vnni); + }, + 32, 32, Choose_UINT8_IP_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + +// See Cosine_INT8_DispatchTable's comment - the same documented skip-alignment special case +// applies here (uint8 cosine's extra norm float shifts effective alignment the same way). +inline constexpr auto Cosine_UINT8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_UINT8_Cosine_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_UINT8_Cosine_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_UINT8_Cosine_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_UINT8_Cosine_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vl && f.avx512vnni); + }, + 32, 0, Choose_UINT8_Cosine_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + } // namespace spaces diff --git a/src/VecSim/spaces/IP_space.cpp b/src/VecSim/spaces/IP_space.cpp index f6dc95f53..d5fc22aef 100644 --- a/src/VecSim/spaces/IP_space.cpp +++ b/src/VecSim/spaces/IP_space.cpp @@ -47,65 +47,16 @@ dist_func_t IP_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_FP32_InnerProduct; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); -#ifdef CPU_FEATURES_ARCH_AARCH64 - -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_FP32_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_FP32_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd) { - return Choose_SQ8_FP32_IP_implementation_NEON(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_SQ8_FP32_DispatchTable); + if (idx == IP_SQ8_FP32_DispatchTable.size()) { + return SQ8_FP32_InnerProduct; } -#endif - -#endif - -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 8 elements (see the residual handling in the kernels). - // Below that, the scalar implementation is at least as fast anyway. - if (dim < 8) { - return ret_dist_func; + const auto &tier = IP_SQ8_FP32_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); } - // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vnni) { - if (dim % 16 == 0) // SQ8 chunk = 16 bytes - *alignment = 16 * sizeof(uint8_t); - return Choose_SQ8_FP32_IP_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#ifdef OPT_AVX2_FMA - if (features.avx2 && features.fma3) { - if (dim % 8 == 0) // SQ8 chunk = 8 bytes - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP32_IP_implementation_AVX2_FMA(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2) { - if (dim % 8 == 0) // SQ8 chunk = 8 bytes - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP32_IP_implementation_AVX2(dim); - } -#endif -#ifdef OPT_SSE4 - if (features.sse4_1) { - if (dim % 4 == 0) // SQ8 chunk = 4 bytes - *alignment = 4 * sizeof(uint8_t); - return Choose_SQ8_FP32_IP_implementation_SSE4(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } // SQ8-FP32: asymmetric cosine distance between SQ8 storage and FP32 query @@ -116,65 +67,16 @@ dist_func_t Cosine_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignm alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_FP32_Cosine; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); -#ifdef CPU_FEATURES_ARCH_AARCH64 - -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_FP32_Cosine_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_FP32_Cosine_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd) { - return Choose_SQ8_FP32_Cosine_implementation_NEON(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, Cosine_SQ8_FP32_DispatchTable); + if (idx == Cosine_SQ8_FP32_DispatchTable.size()) { + return SQ8_FP32_Cosine; } -#endif - -#endif - -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 8 elements (see the residual handling in the kernels). - // Below that, the scalar implementation is at least as fast anyway. - if (dim < 8) { - return ret_dist_func; + const auto &tier = Cosine_SQ8_FP32_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); } - // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vnni) { - if (dim % 16 == 0) // SQ8 chunk = 16 bytes - *alignment = 16 * sizeof(uint8_t); - return Choose_SQ8_FP32_Cosine_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#ifdef OPT_AVX2_FMA - if (features.avx2 && features.fma3) { - if (dim % 8 == 0) // SQ8 chunk = 8 bytes - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP32_Cosine_implementation_AVX2_FMA(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2) { - if (dim % 8 == 0) // SQ8 chunk = 8 bytes - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP32_Cosine_implementation_AVX2(dim); - } -#endif -#ifdef OPT_SSE4 - if (features.sse4_1) { - if (dim % 4 == 0) // SQ8 chunk = 4 bytes - *alignment = 4 * sizeof(uint8_t); - return Choose_SQ8_FP32_Cosine_implementation_SSE4(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } // SQ8-FP16: asymmetric inner product distance between SQ8 storage and FP16 query. @@ -185,73 +87,17 @@ dist_func_t IP_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_FP16_InnerProduct; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 16) { - return ret_dist_func; - } // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. - // AVX-512 tier only needs AVX-512F (cvtph_ps is part of AVX-512F, no VNNI/BW/VL required). -#ifdef OPT_AVX512F - if (features.avx512f) { - if (dim % 16 == 0) // SQ8 chunk = 16 bytes - *alignment = 16 * sizeof(uint8_t); - return Choose_SQ8_FP16_IP_implementation_AVX512F(dim); - } -#endif - // F16C is required by every non-AVX-512 SQ8↔FP16 tier (vcvtph2ps), so the guard is hoisted - // around all three. -#ifdef OPT_F16C -#ifdef OPT_AVX2_FMA - if (features.avx2 && features.fma3 && features.f16c) { - if (dim % 8 == 0) // SQ8 chunk = 8 bytes - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP16_IP_implementation_AVX2_FMA(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2 && features.f16c) { - if (dim % 8 == 0) - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP16_IP_implementation_AVX2(dim); - } -#endif -#ifdef OPT_SSE4 - // F16C is VEX-encoded — require AVX as well, matching the existing F16C/FP16 dispatcher. - if (features.sse4_1 && features.f16c && features.avx) { - if (dim % 4 == 0) - *alignment = 4 * sizeof(uint8_t); - return Choose_SQ8_FP16_IP_implementation_SSE4(dim); - } -#endif -#endif // OPT_F16C -#endif // x86_64 -#ifdef CPU_FEATURES_ARCH_AARCH64 - if (dim < 16) { - return ret_dist_func; - } -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_FP16_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_FP16_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_HP - if (features.asimdfhm) { - return Choose_SQ8_FP16_IP_implementation_NEON_FHM(dim); - } - if (features.asimdhp) { - return Choose_SQ8_FP16_IP_implementation_NEON_HP(dim); - } -#endif -#endif // CPU_FEATURES_ARCH_AARCH64 - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_SQ8_FP16_DispatchTable); + if (idx == IP_SQ8_FP16_DispatchTable.size()) { + return SQ8_FP16_InnerProduct; + } + const auto &tier = IP_SQ8_FP16_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); + } + return tier.chooser(dim); } // SQ8-FP16: asymmetric cosine distance between SQ8 storage and FP16 query. @@ -262,68 +108,16 @@ dist_func_t Cosine_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignm alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_FP16_Cosine; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 16) { - return ret_dist_func; - } -#ifdef OPT_AVX512F - if (features.avx512f) { - if (dim % 16 == 0) - *alignment = 16 * sizeof(uint8_t); - return Choose_SQ8_FP16_Cosine_implementation_AVX512F(dim); - } -#endif -#ifdef OPT_F16C -#ifdef OPT_AVX2_FMA - if (features.avx2 && features.fma3 && features.f16c) { - if (dim % 8 == 0) - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP16_Cosine_implementation_AVX2_FMA(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2 && features.f16c) { - if (dim % 8 == 0) - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP16_Cosine_implementation_AVX2(dim); - } -#endif -#ifdef OPT_SSE4 - if (features.sse4_1 && features.f16c && features.avx) { - if (dim % 4 == 0) - *alignment = 4 * sizeof(uint8_t); - return Choose_SQ8_FP16_Cosine_implementation_SSE4(dim); - } -#endif -#endif // OPT_F16C -#endif // x86_64 -#ifdef CPU_FEATURES_ARCH_AARCH64 - if (dim < 16) { - return ret_dist_func; - } -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_FP16_Cosine_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_FP16_Cosine_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_HP - if (features.asimdfhm) { - return Choose_SQ8_FP16_Cosine_implementation_NEON_FHM(dim); - } - if (features.asimdhp) { - return Choose_SQ8_FP16_Cosine_implementation_NEON_HP(dim); - } -#endif -#endif // CPU_FEATURES_ARCH_AARCH64 - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, Cosine_SQ8_FP16_DispatchTable); + if (idx == Cosine_SQ8_FP16_DispatchTable.size()) { + return SQ8_FP16_Cosine; + } + const auto &tier = Cosine_SQ8_FP16_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); + } + return tier.chooser(dim); } // SQ8-to-SQ8 Inner Product distance function (both vectors are uint8 quantized with precomputed @@ -335,43 +129,16 @@ dist_func_t IP_SQ8_SQ8_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_SQ8_InnerProduct; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_SQ8_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_SQ8_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - if (features.asimddp && dim >= 16) { - return Choose_SQ8_SQ8_IP_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_SQ8_SQ8_IP_implementation_NEON(dim); - } -#endif -#endif // AARCH64 - -#ifdef CPU_FEATURES_ARCH_X86_64 -#ifdef OPT_AVX512_F_BW_VL_VNNI - // AVX512 VNNI SQ8_SQ8 uses 64-element chunks; residual handling is in 32-byte sub-chunks. - if (dim >= 64 && features.avx512f && features.avx512bw && features.avx512vnni) { - if (dim % 32 == 0) // align to 256 bits when there is no offsetting residual - *alignment = 32 * sizeof(uint8_t); - return Choose_SQ8_SQ8_IP_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_SQ8_SQ8_DispatchTable); + if (idx == IP_SQ8_SQ8_DispatchTable.size()) { + return SQ8_SQ8_InnerProduct; + } + const auto &tier = IP_SQ8_SQ8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); + } + return tier.chooser(dim); } // SQ8-to-SQ8 Cosine distance function (both vectors are uint8 quantized with precomputed sum) @@ -382,43 +149,16 @@ dist_func_t Cosine_SQ8_SQ8_GetDistFunc(size_t dim, unsigned char *alignme alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_SQ8_Cosine; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_SQ8_Cosine_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_SQ8_Cosine_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - if (features.asimddp && dim >= 16) { - return Choose_SQ8_SQ8_Cosine_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_SQ8_SQ8_Cosine_implementation_NEON(dim); - } -#endif -#endif // AARCH64 - -#ifdef CPU_FEATURES_ARCH_X86_64 -#ifdef OPT_AVX512_F_BW_VL_VNNI - // AVX512 VNNI SQ8_SQ8 uses 64-element chunks; residual handling is in 32-byte sub-chunks. - if (dim >= 64 && features.avx512f && features.avx512bw && features.avx512vnni) { - if (dim % 32 == 0) // align to 256 bits when there is no offsetting residual - *alignment = 32 * sizeof(uint8_t); - return Choose_SQ8_SQ8_Cosine_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, Cosine_SQ8_SQ8_DispatchTable); + if (idx == Cosine_SQ8_SQ8_DispatchTable.size()) { + return SQ8_SQ8_Cosine; + } + const auto &tier = Cosine_SQ8_SQ8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); + } + return tier.chooser(dim); } dist_func_t IP_FP32_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -446,57 +186,16 @@ dist_func_t IP_FP64_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = FP64_InnerProduct; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_FP64_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_FP64_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd) { - return Choose_FP64_IP_implementation_NEON(dim); - } -#endif - -#endif - -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 4 doubles (see the residual handling in the kernels). - // Below that, the scalar implementation is at least as fast anyway. - if (dim < 4) { - return ret_dist_func; - } -#ifdef OPT_AVX512F - if (features.avx512f) { - if (dim % 8 == 0) // no point in aligning if we have an offsetting residual - *alignment = 8 * sizeof(double); // handles 8 doubles - return Choose_FP64_IP_implementation_AVX512F(dim); - } -#endif -#ifdef OPT_AVX - if (features.avx) { - if (dim % 4 == 0) // no point in aligning if we have an offsetting residual - *alignment = 4 * sizeof(double); // handles 4 doubles - return Choose_FP64_IP_implementation_AVX(dim); - } -#endif -#ifdef OPT_SSE - if (features.sse) { - if (dim % 2 == 0) // no point in aligning if we have an offsetting residual - *alignment = 2 * sizeof(double); // handles 2 doubles - return Choose_FP64_IP_implementation_SSE(dim); - } -#endif -#endif // __x86_64__ */ - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_FP64_DispatchTable); + if (idx == IP_FP64_DispatchTable.size()) { + return FP64_InnerProduct; + } + const auto &tier = IP_FP64_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(double); + } + return tier.chooser(dim); } dist_func_t IP_BF16_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -505,61 +204,21 @@ dist_func_t IP_BF16_GetDistFunc(size_t dim, unsigned char *alignment, con alignment = &dummy_alignment; } - dist_func_t ret_dist_func = BF16_InnerProduct_LittleEndian; + // Big/little-endian is not a tier-selection concern - handled before the table is consulted. if (!is_little_endian()) { return BF16_InnerProduct_BigEndian; } - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#if defined(CPU_FEATURES_ARCH_AARCH64) -#ifdef OPT_SVE_BF16 - if (features.svebf16) { - return Choose_BF16_IP_implementation_SVE_BF16(dim); - } -#endif -#ifdef OPT_NEON_BF16 - if (features.bf16 && dim >= 8) { // Optimization assumes at least 8 BF16s (full chunk) - return Choose_BF16_IP_implementation_NEON_BF16(dim); - } -#endif -#endif // AARCH64 -#if defined(CPU_FEATURES_ARCH_X86_64) - // Optimizations assume at least 32 bfloats. If we have less, we use the naive implementation. - if (dim < 32) { - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_BF16_DispatchTable); + if (idx == IP_BF16_DispatchTable.size()) { + return BF16_InnerProduct_LittleEndian; } - -#ifdef OPT_AVX512_BF16_VL - if (features.avx512_bf16 && features.avx512vl) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(bfloat16); // align to 512 bits. - return Choose_BF16_IP_implementation_AVX512BF16_VL(dim); - } -#endif -#ifdef OPT_AVX512_BW_VBMI2 - if (features.avx512bw && features.avx512vbmi2) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(bfloat16); // align to 512 bits. - return Choose_BF16_IP_implementation_AVX512BW_VBMI2(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2) { - if (dim % 16 == 0) // no point in aligning if we have an offsetting residual - *alignment = 16 * sizeof(bfloat16); // align to 256 bits. - return Choose_BF16_IP_implementation_AVX2(dim); - } -#endif -#ifdef OPT_SSE3 - if (features.sse3) { - if (dim % 8 == 0) // no point in aligning if we have an offsetting residual - *alignment = 8 * sizeof(bfloat16); // align to 128 bits. - return Choose_BF16_IP_implementation_SSE3(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + const auto &tier = IP_BF16_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(bfloat16); + } + return tier.chooser(dim); } dist_func_t IP_FP16_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -567,57 +226,17 @@ dist_func_t IP_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con if (alignment == nullptr) { alignment = &dummy_alignment; } - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - - dist_func_t ret_dist_func = FP16_InnerProduct; -#if defined(CPU_FEATURES_ARCH_AARCH64) -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_FP16_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_FP16_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_HP - if (features.asimdhp && dim >= 8) { // Optimization assumes at least 8 16FPs (full chunk) - return Choose_FP16_IP_implementation_NEON_HP(dim); - } -#endif -#endif - -#if defined(CPU_FEATURES_ARCH_X86_64) - // Each tier has a minimal dimension implied by its residual handling: the AVX512FP16_VL - // kernel loads full 512-bit blocks (32 elements), the AVX512F kernel loads full 256-bit - // blocks (16 elements), and the F16C kernel loads full 128-bit blocks (8 elements). -#ifdef OPT_AVX512_FP16_VL - // More details about the dimension limitation can be found in this PR's description: - // https://github.com/RedisAI/VectorSimilarity/pull/477 - if (dim >= 32 && features.avx512_fp16 && features.avx512vl) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(float16); // handles 32 floats - return Choose_FP16_IP_implementation_AVX512FP16_VL(dim); - } -#endif -#ifdef OPT_AVX512F - if (dim >= 16 && features.avx512f) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(float16); // handles 32 floats - return Choose_FP16_IP_implementation_AVX512F(dim); - } -#endif -#ifdef OPT_F16C - if (dim >= 8 && features.f16c && features.fma3 && features.avx) { - if (dim % 16 == 0) // no point in aligning if we have an offsetting residual - *alignment = 16 * sizeof(float16); // handles 16 floats - return Choose_FP16_IP_implementation_F16C(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_FP16_DispatchTable); + if (idx == IP_FP16_DispatchTable.size()) { + return FP16_InnerProduct; + } + const auto &tier = IP_FP16_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(float16); + } + return tier.chooser(dim); } dist_func_t IP_INT8_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -626,47 +245,16 @@ dist_func_t IP_INT8_GetDistFunc(size_t dim, unsigned char *alignment, con alignment = &dummy_alignment; } - dist_func_t ret_dist_func = INT8_InnerProduct; - - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_INT8_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_INT8_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD // Should be the first check, as it is the most optimized - if (features.asimddp && dim >= 16) { - return Choose_INT8_IP_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_INT8_IP_implementation_NEON(dim); - } -#endif -#endif -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 32 int8. If we have less, we use the naive implementation. - if (dim < 32) { - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_INT8_DispatchTable); + if (idx == IP_INT8_DispatchTable.size()) { + return INT8_InnerProduct; } - -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vl && features.avx512vnni) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(int8_t); // align to 256 bits. - return Choose_INT8_IP_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + const auto &tier = IP_INT8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(int8_t); + } + return tier.chooser(dim); } dist_func_t Cosine_INT8_GetDistFunc(size_t dim, unsigned char *alignment, @@ -676,49 +264,16 @@ dist_func_t Cosine_INT8_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = INT8_Cosine; - - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_INT8_Cosine_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_INT8_Cosine_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - if (features.asimddp && dim >= 16) { - return Choose_INT8_Cosine_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_INT8_Cosine_implementation_NEON(dim); - } -#endif -#endif -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 32) { - return ret_dist_func; - } -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vl && features.avx512vnni) { - // For int8 vectors with cosine distance, the extra float for the norm shifts alignment to - // `(dim + sizeof(float)) % 32`. - // Vectors satisfying this have a residual, causing offset loads during calculation. - // To avoid complexity, we skip alignment here, assuming the performance impact is - // negligible. - return Choose_INT8_Cosine_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif - -#endif // __x86_64__ - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, Cosine_INT8_DispatchTable); + if (idx == Cosine_INT8_DispatchTable.size()) { + return INT8_Cosine; + } + const auto &tier = Cosine_INT8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(int8_t); + } + return tier.chooser(dim); } dist_func_t IP_UINT8_GetDistFunc(size_t dim, unsigned char *alignment, @@ -728,45 +283,16 @@ dist_func_t IP_UINT8_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = UINT8_InnerProduct; - - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_UINT8_IP_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_UINT8_IP_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - if (features.asimddp && dim >= 16) { - return Choose_UINT8_IP_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_UINT8_IP_implementation_NEON(dim); - } -#endif -#endif -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 32) { - return ret_dist_func; - } -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vl && features.avx512vnni) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(uint8_t); // align to 256 bits. - return Choose_UINT8_IP_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, IP_UINT8_DispatchTable); + if (idx == IP_UINT8_DispatchTable.size()) { + return UINT8_InnerProduct; + } + const auto &tier = IP_UINT8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); + } + return tier.chooser(dim); } dist_func_t Cosine_UINT8_GetDistFunc(size_t dim, unsigned char *alignment, @@ -776,48 +302,16 @@ dist_func_t Cosine_UINT8_GetDistFunc(size_t dim, unsigned char *alignment alignment = &dummy_alignment; } - dist_func_t ret_dist_func = UINT8_Cosine; - - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_UINT8_Cosine_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_UINT8_Cosine_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - if (features.asimddp && dim >= 16) { - return Choose_UINT8_Cosine_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_UINT8_Cosine_implementation_NEON(dim); - } -#endif -#endif -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 32) { - return ret_dist_func; - } -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vl && features.avx512vnni) { - // For uint8 vectors with cosine distance, the extra float for the norm shifts alignment to - // `(dim + sizeof(float)) % 32`. - // Vectors satisfying this have a residual, causing offset loads during calculation. - // To avoid complexity, we skip alignment here, assuming the performance impact is - // negligible. - return Choose_UINT8_Cosine_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#endif // __x86_64__ - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, Cosine_UINT8_DispatchTable); + if (idx == Cosine_UINT8_DispatchTable.size()) { + return UINT8_Cosine; + } + const auto &tier = Cosine_UINT8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); + } + return tier.chooser(dim); } } // namespace spaces diff --git a/src/VecSim/spaces/L2_dispatch_tables.h b/src/VecSim/spaces/L2_dispatch_tables.h index c08489ab8..a8660ff4d 100644 --- a/src/VecSim/spaces/L2_dispatch_tables.h +++ b/src/VecSim/spaces/L2_dispatch_tables.h @@ -11,14 +11,32 @@ #include "VecSim/spaces/dispatch_tier.h" #include "VecSim/spaces/functions/SVE2.h" #include "VecSim/spaces/functions/SVE.h" +#include "VecSim/spaces/functions/SVE_BF16.h" #include "VecSim/spaces/functions/NEON.h" +#include "VecSim/spaces/functions/NEON_DOTPROD.h" +#include "VecSim/spaces/functions/NEON_HP.h" +#include "VecSim/spaces/functions/NEON_BF16.h" #include "VecSim/spaces/functions/AVX512F.h" +#include "VecSim/spaces/functions/AVX512BW_VBMI2.h" +#include "VecSim/spaces/functions/AVX512FP16_VL.h" +#include "VecSim/spaces/functions/AVX512F_BW_VL_VNNI.h" #include "VecSim/spaces/functions/AVX.h" +#include "VecSim/spaces/functions/AVX2.h" +#include "VecSim/spaces/functions/AVX2_F16C.h" +#include "VecSim/spaces/functions/AVX2_FMA.h" +#include "VecSim/spaces/functions/AVX2_FMA_F16C.h" +#include "VecSim/spaces/functions/F16C.h" #include "VecSim/spaces/functions/SSE.h" +#include "VecSim/spaces/functions/SSE3.h" +#include "VecSim/spaces/functions/SSE4.h" +#include "VecSim/spaces/functions/SSE4_F16C.h" /* * `inline constexpr` dispatch tables for L2_space.cpp's GetDistFunc-per-combo functions - see - * IP_dispatch_tables.h for the full rationale (same reasons apply here, mirrored for L2). + * IP_dispatch_tables.h for the full rationale (same reasons apply here, mirrored for L2). Note + * L2 has no Cosine tables at all: for int8/uint8, Cosine is dispatched through IP_space.cpp's + * Cosine_* functions, never through L2_space.cpp - so L2's AVX-512 int8/uint8 tiers set alignment + * normally, with none of IP's documented cosine skip-alignment special case. */ namespace spaces { @@ -53,4 +71,241 @@ inline constexpr auto L2_FP32_DispatchTable = std::array{ #endif }; +// Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. +inline constexpr auto L2_SQ8_FP32_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_SQ8_FP32_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_SQ8_FP32_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 0, 0, + Choose_SQ8_FP32_L2_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vnni); + }, + 8, 16, Choose_SQ8_FP32_L2_implementation_AVX512F_BW_VL_VNNI}, +#endif +#ifdef OPT_AVX2_FMA + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.fma3); }, 8, + 8, Choose_SQ8_FP32_L2_implementation_AVX2_FMA}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx2; }, 8, 8, + Choose_SQ8_FP32_L2_implementation_AVX2}, +#endif +#ifdef OPT_SSE4 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse4_1; }, 8, 4, + Choose_SQ8_FP32_L2_implementation_SSE4}, +#endif +}; + +// min_dim=16 on every row - both the x86 and the ARM block in the original cascade shared one +// `if (dim < 16) return ret_dist_func;` gate over all their rows. +inline constexpr auto L2_SQ8_FP16_DispatchTable = std::array{ +#ifdef OPT_AVX512F + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 16, 16, + Choose_SQ8_FP16_L2_implementation_AVX512F}, +#endif +#ifdef OPT_F16C +#ifdef OPT_AVX2_FMA + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.fma3 && f.f16c); }, 16, 8, + Choose_SQ8_FP16_L2_implementation_AVX2_FMA}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)(f.avx2 && f.f16c); }, 16, + 8, Choose_SQ8_FP16_L2_implementation_AVX2}, +#endif +#ifdef OPT_SSE4 + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.sse4_1 && f.f16c && f.avx); }, 16, 4, + Choose_SQ8_FP16_L2_implementation_SSE4}, +#endif +#endif // OPT_F16C +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 16, 0, + Choose_SQ8_FP16_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 16, 0, + Choose_SQ8_FP16_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON_HP + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdfhm; }, 16, 0, + Choose_SQ8_FP16_L2_implementation_NEON_FHM}, + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdhp; }, 16, 0, + Choose_SQ8_FP16_L2_implementation_NEON_HP}, +#endif +}; + +inline constexpr auto L2_FP64_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_FP64_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_FP64_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 0, 0, + Choose_FP64_L2_implementation_NEON}, +#endif +#ifdef OPT_AVX512F + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 4, 8, + Choose_FP64_L2_implementation_AVX512F}, +#endif +#ifdef OPT_AVX + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx; }, 4, 4, + Choose_FP64_L2_implementation_AVX}, +#endif +#ifdef OPT_SSE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse; }, 4, 2, + Choose_FP64_L2_implementation_SSE}, +#endif +}; + +// The BF16 big/little-endian split stays as a special-cased branch in L2_BF16_GetDistFunc itself, +// before this table is ever consulted - see L2_space.cpp (mirrors IP_BF16_DispatchTable). +inline constexpr auto L2_BF16_DispatchTable = std::array{ +#ifdef OPT_SVE_BF16 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.svebf16; }, 0, 0, + Choose_BF16_L2_implementation_SVE_BF16}, +#endif +#ifdef OPT_NEON_BF16 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.bf16; }, 8, 0, + Choose_BF16_L2_implementation_NEON_BF16}, +#endif +#ifdef OPT_AVX512_BW_VBMI2 + // Note: unlike IP_BF16_DispatchTable, there is no AVX512BF16_VL row here - the original + // L2_BF16_GetDistFunc cascade never had one, only AVX512BW_VBMI2. + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx512bw && f.avx512vbmi2); }, 32, 32, + Choose_BF16_L2_implementation_AVX512BW_VBMI2}, +#endif +#ifdef OPT_AVX2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx2; }, 32, 16, + Choose_BF16_L2_implementation_AVX2}, +#endif +#ifdef OPT_SSE3 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sse3; }, 32, 8, + Choose_BF16_L2_implementation_SSE3}, +#endif +}; + +inline constexpr auto L2_FP16_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_FP16_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_FP16_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON_HP + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimdhp; }, 8, 0, + Choose_FP16_L2_implementation_NEON_HP}, +#endif +#ifdef OPT_AVX512_FP16_VL + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.avx512_fp16 && f.avx512vl); }, 32, 32, + Choose_FP16_L2_implementation_AVX512FP16_VL}, +#endif +#ifdef OPT_AVX512F + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.avx512f; }, 16, 32, + Choose_FP16_L2_implementation_AVX512F}, +#endif +#ifdef OPT_F16C + DispatchTier{ + [](const FeaturesType &f) noexcept { return (bool)(f.f16c && f.fma3 && f.avx); }, 8, 16, + Choose_FP16_L2_implementation_F16C}, +#endif +}; + +// Unlike IP_space.cpp's Cosine_INT8/Cosine_UINT8, there is no skip-alignment special case here - +// L2's AVX-512 tier sets alignment normally. +inline constexpr auto L2_INT8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_INT8_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_INT8_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_INT8_L2_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_INT8_L2_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vl && f.avx512vnni); + }, + 32, 32, Choose_INT8_L2_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + +inline constexpr auto L2_UINT8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_UINT8_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_UINT8_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_UINT8_L2_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_UINT8_L2_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vl && f.avx512vnni); + }, + 32, 32, Choose_UINT8_L2_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + +// Both operands are SQ8 with a precomputed sum; the ARM rows carry their own per-row dim>=16 +// floor, and the single x86 tier uses 64-element chunks with its own dim>=64 floor - mirrors +// IP_SQ8_SQ8_DispatchTable. +inline constexpr auto L2_SQ8_SQ8_DispatchTable = std::array{ +#ifdef OPT_SVE2 + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve2; }, 0, 0, + Choose_SQ8_SQ8_L2_implementation_SVE2}, +#endif +#ifdef OPT_SVE + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.sve; }, 0, 0, + Choose_SQ8_SQ8_L2_implementation_SVE}, +#endif +#ifdef OPT_NEON_DOTPROD + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimddp; }, 16, 0, + Choose_SQ8_SQ8_L2_implementation_NEON_DOTPROD}, +#endif +#ifdef OPT_NEON + DispatchTier{[](const FeaturesType &f) noexcept { return (bool)f.asimd; }, 16, 0, + Choose_SQ8_SQ8_L2_implementation_NEON}, +#endif +#ifdef OPT_AVX512_F_BW_VL_VNNI + DispatchTier{[](const FeaturesType &f) noexcept { + return (bool)(f.avx512f && f.avx512bw && f.avx512vnni); + }, + 64, 32, Choose_SQ8_SQ8_L2_implementation_AVX512F_BW_VL_VNNI}, +#endif +}; + } // namespace spaces diff --git a/src/VecSim/spaces/L2_space.cpp b/src/VecSim/spaces/L2_space.cpp index d048a0c99..5c50f1708 100644 --- a/src/VecSim/spaces/L2_space.cpp +++ b/src/VecSim/spaces/L2_space.cpp @@ -47,64 +47,16 @@ dist_func_t L2_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_FP32_L2Sqr; - - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_FP32_L2_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_FP32_L2_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd) { - return Choose_SQ8_FP32_L2_implementation_NEON(dim); - } -#endif -#endif - -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 8 elements (see the residual handling in the kernels). - // Below that, the scalar implementation is at least as fast anyway. - if (dim < 8) { - return ret_dist_func; - } - // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vnni) { - if (dim % 16 == 0) // SQ8 chunk = 16 bytes; no point in aligning if there's a residual - *alignment = 16 * sizeof(uint8_t); - return Choose_SQ8_FP32_L2_implementation_AVX512F_BW_VL_VNNI(dim); - } -#endif -#ifdef OPT_AVX2_FMA - if (features.avx2 && features.fma3) { - if (dim % 8 == 0) // SQ8 chunk = 8 bytes - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP32_L2_implementation_AVX2_FMA(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2) { - if (dim % 8 == 0) // SQ8 chunk = 8 bytes - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP32_L2_implementation_AVX2(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_SQ8_FP32_DispatchTable); + if (idx == L2_SQ8_FP32_DispatchTable.size()) { + return SQ8_FP32_L2Sqr; } -#endif -#ifdef OPT_SSE4 - if (features.sse4_1) { - if (dim % 4 == 0) // SQ8 chunk = 4 bytes - *alignment = 4 * sizeof(uint8_t); - return Choose_SQ8_FP32_L2_implementation_SSE4(dim); + const auto &tier = L2_SQ8_FP32_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } // SQ8-FP16: asymmetric L2 distance between SQ8 storage and FP16 query. @@ -115,72 +67,16 @@ dist_func_t L2_SQ8_FP16_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_FP16_L2Sqr; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 16) { - return ret_dist_func; - } - // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract. - // AVX-512 tier only needs AVX-512F (cvtph_ps is part of AVX-512F, no VNNI/BW/VL required). -#ifdef OPT_AVX512F - if (features.avx512f) { - if (dim % 16 == 0) - *alignment = 16 * sizeof(uint8_t); - return Choose_SQ8_FP16_L2_implementation_AVX512F(dim); - } -#endif - // F16C is required by every non-AVX-512 SQ8↔FP16 tier (vcvtph2ps), so the guard is hoisted - // around all three. -#ifdef OPT_F16C -#ifdef OPT_AVX2_FMA - if (features.avx2 && features.fma3 && features.f16c) { - if (dim % 8 == 0) - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP16_L2_implementation_AVX2_FMA(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2 && features.f16c) { - if (dim % 8 == 0) - *alignment = 8 * sizeof(uint8_t); - return Choose_SQ8_FP16_L2_implementation_AVX2(dim); - } -#endif -#ifdef OPT_SSE4 - if (features.sse4_1 && features.f16c && features.avx) { - if (dim % 4 == 0) - *alignment = 4 * sizeof(uint8_t); - return Choose_SQ8_FP16_L2_implementation_SSE4(dim); - } -#endif -#endif // OPT_F16C -#endif // x86_64 -#ifdef CPU_FEATURES_ARCH_AARCH64 - if (dim < 16) { - return ret_dist_func; - } -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_FP16_L2_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_FP16_L2_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_HP - if (features.asimdfhm) { - return Choose_SQ8_FP16_L2_implementation_NEON_FHM(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_SQ8_FP16_DispatchTable); + if (idx == L2_SQ8_FP16_DispatchTable.size()) { + return SQ8_FP16_L2Sqr; } - if (features.asimdhp) { - return Choose_SQ8_FP16_L2_implementation_NEON_HP(dim); + const auto &tier = L2_SQ8_FP16_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); } -#endif -#endif // CPU_FEATURES_ARCH_AARCH64 - return ret_dist_func; + return tier.chooser(dim); } dist_func_t L2_FP32_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -208,56 +104,16 @@ dist_func_t L2_FP64_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = FP64_L2Sqr; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_FP64_L2_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_FP64_L2_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd) { - return Choose_FP64_L2_implementation_NEON(dim); - } -#endif -#endif - -#ifdef CPU_FEATURES_ARCH_X86_64 - // Optimizations assume at least 4 doubles (see the residual handling in the kernels). - // Below that, the scalar implementation is at least as fast anyway. - if (dim < 4) { - return ret_dist_func; - } -#ifdef OPT_AVX512F - if (features.avx512f) { - if (dim % 8 == 0) // no point in aligning if we have an offsetting residual - *alignment = 8 * sizeof(double); // handles 8 doubles - return Choose_FP64_L2_implementation_AVX512F(dim); - } -#endif -#ifdef OPT_AVX - if (features.avx) { - if (dim % 4 == 0) // no point in aligning if we have an offsetting residual - *alignment = 4 * sizeof(double); // handles 4 doubles - return Choose_FP64_L2_implementation_AVX(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_FP64_DispatchTable); + if (idx == L2_FP64_DispatchTable.size()) { + return FP64_L2Sqr; } -#endif -#ifdef OPT_SSE - if (features.sse) { - if (dim % 2 == 0) // no point in aligning if we have an offsetting residual - *alignment = 2 * sizeof(double); // handles 2 doubles - return Choose_FP64_L2_implementation_SSE(dim); + const auto &tier = L2_FP64_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(double); } -#endif -#endif // __x86_64__ */ - return ret_dist_func; + return tier.chooser(dim); } dist_func_t L2_BF16_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -266,53 +122,21 @@ dist_func_t L2_BF16_GetDistFunc(size_t dim, unsigned char *alignment, con alignment = &dummy_alignment; } - dist_func_t ret_dist_func = BF16_L2Sqr_LittleEndian; + // Big/little-endian is not a tier-selection concern - handled before the table is consulted. if (!is_little_endian()) { return BF16_L2Sqr_BigEndian; } - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); -#if defined(CPU_FEATURES_ARCH_AARCH64) -#ifdef OPT_SVE_BF16 - if (features.svebf16) { - return Choose_BF16_L2_implementation_SVE_BF16(dim); - } -#endif -#ifdef OPT_NEON_BF16 - if (features.bf16 && dim >= 8) { // Optimization assumes at least 8 BF16s (full chunk) - return Choose_BF16_L2_implementation_NEON_BF16(dim); - } -#endif -#endif // AARCH64 - -#if defined(CPU_FEATURES_ARCH_X86_64) - // Optimizations assume at least 32 bfloats. If we have less, we use the naive implementation. - if (dim < 32) { - return ret_dist_func; - } -#ifdef OPT_AVX512_BW_VBMI2 - if (features.avx512bw && features.avx512vbmi2) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(bfloat16); // align to 512 bits. - return Choose_BF16_L2_implementation_AVX512BW_VBMI2(dim); - } -#endif -#ifdef OPT_AVX2 - if (features.avx2) { - if (dim % 16 == 0) // no point in aligning if we have an offsetting residual - *alignment = 16 * sizeof(bfloat16); // align to 256 bits. - return Choose_BF16_L2_implementation_AVX2(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_BF16_DispatchTable); + if (idx == L2_BF16_DispatchTable.size()) { + return BF16_L2Sqr_LittleEndian; } -#endif -#ifdef OPT_SSE3 - if (features.sse3) { - if (dim % 8 == 0) // no point in aligning if we have an offsetting residual - *alignment = 8 * sizeof(bfloat16); // align to 128 bits. - return Choose_BF16_L2_implementation_SSE3(dim); + const auto &tier = L2_BF16_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(bfloat16); } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } dist_func_t L2_FP16_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -320,57 +144,17 @@ dist_func_t L2_FP16_GetDistFunc(size_t dim, unsigned char *alignment, con if (alignment == nullptr) { alignment = &dummy_alignment; } - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - - dist_func_t ret_dist_func = FP16_L2Sqr; - -#if defined(CPU_FEATURES_ARCH_AARCH64) -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_FP16_L2_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_FP16_L2_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_HP - if (features.asimdhp && dim >= 8) { // Optimization assumes at least 8 16FPs (full chunk) - return Choose_FP16_L2_implementation_NEON_HP(dim); - } -#endif -#endif // CPU_FEATURES_ARCH_AARCH64 -#if defined(CPU_FEATURES_ARCH_X86_64) - // Each tier has a minimal dimension implied by its residual handling: the AVX512FP16_VL - // kernel loads full 512-bit blocks (32 elements), the AVX512F kernel loads full 256-bit - // blocks (16 elements), and the F16C kernel loads full 128-bit blocks (8 elements). -#ifdef OPT_AVX512_FP16_VL - // More details about the dimension limitation can be found in this PR's description: - // https://github.com/RedisAI/VectorSimilarity/pull/477 - if (dim >= 32 && features.avx512_fp16 && features.avx512vl) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(float16); // handles 32 floats - return Choose_FP16_L2_implementation_AVX512FP16_VL(dim); - } -#endif -#ifdef OPT_AVX512F - if (dim >= 16 && features.avx512f) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(float16); // handles 32 floats - return Choose_FP16_L2_implementation_AVX512F(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_FP16_DispatchTable); + if (idx == L2_FP16_DispatchTable.size()) { + return FP16_L2Sqr; } -#endif -#ifdef OPT_F16C - if (dim >= 8 && features.f16c && features.fma3 && features.avx) { - if (dim % 16 == 0) // no point in aligning if we have an offsetting residual - *alignment = 16 * sizeof(float16); // handles 16 floats - return Choose_FP16_L2_implementation_F16C(dim); + const auto &tier = L2_FP16_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(float16); } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } dist_func_t L2_INT8_GetDistFunc(size_t dim, unsigned char *alignment, const void *arch_opt) { @@ -379,45 +163,16 @@ dist_func_t L2_INT8_GetDistFunc(size_t dim, unsigned char *alignment, con alignment = &dummy_alignment; } - dist_func_t ret_dist_func = INT8_L2Sqr; - - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_INT8_L2_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_INT8_L2_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - if (features.asimddp && dim >= 16) { - return Choose_INT8_L2_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_INT8_L2_implementation_NEON(dim); - } -#endif -#endif -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 32) { - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_INT8_DispatchTable); + if (idx == L2_INT8_DispatchTable.size()) { + return INT8_L2Sqr; } -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vl && features.avx512vnni) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(int8_t); // align to 256 bits. - return Choose_INT8_L2_implementation_AVX512F_BW_VL_VNNI(dim); + const auto &tier = L2_INT8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(int8_t); } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } dist_func_t L2_UINT8_GetDistFunc(size_t dim, unsigned char *alignment, @@ -427,45 +182,16 @@ dist_func_t L2_UINT8_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = UINT8_L2Sqr; - // Optimizations assume at least 32 uint8. If we have less, we use the naive implementation. - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_UINT8_L2_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_UINT8_L2_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - if (features.asimddp && dim >= 16) { - return Choose_UINT8_L2_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (features.asimd && dim >= 16) { - return Choose_UINT8_L2_implementation_NEON(dim); - } -#endif -#endif // __aarch64__ -#ifdef CPU_FEATURES_ARCH_X86_64 - if (dim < 32) { - return ret_dist_func; + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_UINT8_DispatchTable); + if (idx == L2_UINT8_DispatchTable.size()) { + return UINT8_L2Sqr; } -#ifdef OPT_AVX512_F_BW_VL_VNNI - if (features.avx512f && features.avx512bw && features.avx512vl && features.avx512vnni) { - if (dim % 32 == 0) // no point in aligning if we have an offsetting residual - *alignment = 32 * sizeof(int8_t); // align to 256 bits. - return Choose_UINT8_L2_implementation_AVX512F_BW_VL_VNNI(dim); + const auto &tier = L2_UINT8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } // SQ8-to-SQ8 L2 squared distance function (both vectors are uint8 quantized) @@ -476,44 +202,16 @@ dist_func_t L2_SQ8_SQ8_GetDistFunc(size_t dim, unsigned char *alignment, alignment = &dummy_alignment; } - dist_func_t ret_dist_func = SQ8_SQ8_L2Sqr; - [[maybe_unused]] auto features = getCpuOptimizationFeatures(arch_opt); - -#ifdef CPU_FEATURES_ARCH_AARCH64 -#ifdef OPT_SVE2 - if (features.sve2) { - return Choose_SQ8_SQ8_L2_implementation_SVE2(dim); - } -#endif -#ifdef OPT_SVE - if (features.sve) { - return Choose_SQ8_SQ8_L2_implementation_SVE(dim); - } -#endif -#ifdef OPT_NEON_DOTPROD - // DOTPROD uses integer arithmetic - much faster than float-based NEON - if (dim >= 16 && features.asimddp) { - return Choose_SQ8_SQ8_L2_implementation_NEON_DOTPROD(dim); - } -#endif -#ifdef OPT_NEON - if (dim >= 16 && features.asimd) { - return Choose_SQ8_SQ8_L2_implementation_NEON(dim); + auto features = getCpuOptimizationFeatures(arch_opt); + size_t idx = select_tier_index(features, dim, L2_SQ8_SQ8_DispatchTable); + if (idx == L2_SQ8_SQ8_DispatchTable.size()) { + return SQ8_SQ8_L2Sqr; } -#endif -#endif // AARCH64 - -#ifdef CPU_FEATURES_ARCH_X86_64 -#ifdef OPT_AVX512_F_BW_VL_VNNI - // AVX512 VNNI SQ8_SQ8 uses 64-element chunks; residual handling is in 32-byte sub-chunks. - if (dim >= 64 && features.avx512f && features.avx512bw && features.avx512vnni) { - if (dim % 32 == 0) // align to 256 bits when there is no offsetting residual - *alignment = 32 * sizeof(uint8_t); - return Choose_SQ8_SQ8_L2_implementation_AVX512F_BW_VL_VNNI(dim); + const auto &tier = L2_SQ8_SQ8_DispatchTable[idx]; + if (tier.alignment_chunk_elems != 0 && dim % tier.alignment_chunk_elems == 0) { + *alignment = tier.alignment_chunk_elems * sizeof(uint8_t); } -#endif -#endif // __x86_64__ - return ret_dist_func; + return tier.chooser(dim); } } // namespace spaces diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 2056e393b..b0011fafd 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -988,6 +988,77 @@ TEST(SpacesDispatchTierTest, FP32IPTableZeroMinDimAcceptsSmallestDims) { } #endif // CPU_FEATURES_ARCH_AARCH64 +#ifdef CPU_FEATURES_ARCH_X86_64 +#ifdef OPT_AVX512_F_BW_VL_VNNI +// Cosine's AVX-512 tier deliberately never sets an alignment hint (documented special case - the +// extra norm float shifts effective alignment in a way the original cascade skips computing). +// There is no L2 mirror: L2_INT8/L2_UINT8 set alignment normally. This asymmetry must survive. +TEST(SpacesDispatchTierTest, CosineInt8Avx512SkipsAlignmentButL2DoesNot) { + spaces::FeaturesType avx512{}; + avx512.avx512f = 1; + avx512.avx512bw = 1; + avx512.avx512vl = 1; + avx512.avx512vnni = 1; + + size_t cosine_idx = spaces::select_tier_index(avx512, 1000, spaces::Cosine_INT8_DispatchTable); + ASSERT_NE(cosine_idx, spaces::Cosine_INT8_DispatchTable.size()); + EXPECT_EQ(spaces::Cosine_INT8_DispatchTable[cosine_idx].alignment_chunk_elems, 0u) + << "Cosine_INT8's AVX512 tier must not carry an alignment hint (documented skip case)"; + + size_t l2_idx = spaces::select_tier_index(avx512, 1000, spaces::L2_INT8_DispatchTable); + ASSERT_NE(l2_idx, spaces::L2_INT8_DispatchTable.size()); + EXPECT_EQ(spaces::L2_INT8_DispatchTable[l2_idx].alignment_chunk_elems, 32u) + << "L2_INT8's AVX512 tier sets alignment normally - no skip case here, unlike Cosine"; +} + +TEST(SpacesDispatchTierTest, CosineUint8Avx512SkipsAlignmentButL2DoesNot) { + spaces::FeaturesType avx512{}; + avx512.avx512f = 1; + avx512.avx512bw = 1; + avx512.avx512vl = 1; + avx512.avx512vnni = 1; + + size_t cosine_idx = spaces::select_tier_index(avx512, 1000, spaces::Cosine_UINT8_DispatchTable); + ASSERT_NE(cosine_idx, spaces::Cosine_UINT8_DispatchTable.size()); + EXPECT_EQ(spaces::Cosine_UINT8_DispatchTable[cosine_idx].alignment_chunk_elems, 0u); + + size_t l2_idx = spaces::select_tier_index(avx512, 1000, spaces::L2_UINT8_DispatchTable); + ASSERT_NE(l2_idx, spaces::L2_UINT8_DispatchTable.size()); + EXPECT_EQ(spaces::L2_UINT8_DispatchTable[l2_idx].alignment_chunk_elems, 32u); +} +#endif // OPT_AVX512_F_BW_VL_VNNI + +#ifdef OPT_AVX2 +// IP_BF16 has an AVX512BF16_VL tier that L2_BF16 never had in the original cascade (confirmed by +// re-reading L2_space.cpp - only AVX512BW_VBMI2 exists there). With avx512_bf16 set but no other +// AVX-512 flags, IP must pick the dedicated AVX512BF16_VL tier while L2 must skip straight past +// it to AVX2 (since L2 has no row whose predicate can match avx512_bf16 alone). +TEST(SpacesDispatchTierTest, IPBf16HasAvx512Bf16VlTierThatL2Lacks) { + spaces::FeaturesType avx512_bf16_and_avx2{}; + avx512_bf16_and_avx2.avx512_bf16 = 1; + avx512_bf16_and_avx2.avx512vl = 1; + avx512_bf16_and_avx2.avx2 = 1; + + size_t ip_idx = + spaces::select_tier_index(avx512_bf16_and_avx2, 1000, spaces::IP_BF16_DispatchTable); + ASSERT_NE(ip_idx, spaces::IP_BF16_DispatchTable.size()); +#ifdef OPT_AVX512_BF16_VL + EXPECT_EQ(spaces::IP_BF16_DispatchTable[ip_idx].chooser, + spaces::Choose_BF16_IP_implementation_AVX512BF16_VL); +#endif + + size_t l2_idx = + spaces::select_tier_index(avx512_bf16_and_avx2, 1000, spaces::L2_BF16_DispatchTable); + ASSERT_NE(l2_idx, spaces::L2_BF16_DispatchTable.size()); + EXPECT_EQ(spaces::L2_BF16_DispatchTable[l2_idx].chooser, + spaces::Choose_BF16_L2_implementation_AVX2) + << "L2_BF16 has no AVX512BF16_VL row, so avx512_bf16-only features must fall through to " + "AVX2 (avx512bw/avx512vbmi2 are unset here, so its AVX512BW_VBMI2 row can't match " + "either)"; +} +#endif // OPT_AVX2 +#endif // CPU_FEATURES_ARCH_X86_64 + class FP64SpacesOptimizationTest : public testing::TestWithParam {}; TEST_P(FP64SpacesOptimizationTest, FP64L2SqrTest) {