diff --git a/cmake/dispatch-checks/check-dispatch-execution.cmake b/cmake/dispatch-checks/check-dispatch-execution.cmake index e15eb9f2..f16ff3a2 100644 --- a/cmake/dispatch-checks/check-dispatch-execution.cmake +++ b/cmake/dispatch-checks/check-dispatch-execution.cmake @@ -23,7 +23,7 @@ ##### Everything else about the surface is a property of the symbol table, which a ##### specialization can satisfy while never running: one that disappears behind an ##### `#if` still links and still counts. This breaks on every level's kernel for one -##### extent and reports the one the host's run reaches. +##### extent and type pair, and reports the one the host's run reaches. ##### ##### Only the level this host satisfies is checked. The weaker levels are checked by ##### hosts that satisfy only those, which is what the CI matrix is for. @@ -70,8 +70,12 @@ endif() ##### The candidate symbols ##### -# L2 at float/float (`ff`) for the extent the probe reports: one symbol per level, -# including any AVX_AVAILABILITY::NONE fallback the consumer instantiated itself. +# The type pair the probe calls, Itanium-mangled. It must be one with a kernel at +# every level in the surface, or the entry points legitimately route lower. +set(svs_probe_pair "aa") + +# L2 at that pair for the extent the probe reports: one symbol per level, including +# any AVX_AVAILABILITY::NONE fallback the consumer instantiated itself. execute_process( COMMAND "${SVS_NM}" --defined-only "${SVS_PROBE}" OUTPUT_VARIABLE raw @@ -87,7 +91,7 @@ string(REPLACE "\n" ";" lines "${raw}") foreach(line IN LISTS lines) # The mangled name is the last whitespace-separated field. string(REGEX MATCH "[^ \t]+$" symbol "${line}") - if(NOT symbol MATCHES "^_ZN3svs8distance6L2ImplILm${extent}Eff.*7computeE") + if(NOT symbol MATCHES "^_ZN3svs8distance6L2ImplILm${extent}E${svs_probe_pair}.*7computeE") continue() endif() # Skip GCC's `.isra` clones: gdb reads a dot in a linespec as a file name. @@ -101,10 +105,10 @@ list(REMOVE_DUPLICATES candidates) list(LENGTH candidates n_candidates) if(n_candidates EQUAL 0) message(FATAL_ERROR - "No L2 kernel symbol for extent ${extent} at float/float is defined in " - "${SVS_PROBE}. Either the entry points inlined every kernel, in which case " - "the `extern template` declarations are not in effect, or the surface no " - "longer covers that extent." + "No L2 kernel symbol for extent ${extent} at the mangled type pair " + "'${svs_probe_pair}' is defined in ${SVS_PROBE}. Either the entry points " + "inlined every kernel, in which case the `extern template` declarations are " + "not in effect, or the probe no longer calls that extent and pair." ) endif() diff --git a/cmake/dispatch-checks/check-dispatch-instructions.cmake b/cmake/dispatch-checks/check-dispatch-instructions.cmake index fafd97d2..45e1f069 100644 --- a/cmake/dispatch-checks/check-dispatch-instructions.cmake +++ b/cmake/dispatch-checks/check-dispatch-instructions.cmake @@ -55,7 +55,7 @@ set(svs_budget_table "x86-64||ymm zmm mask vnni" "haswell|ymm|zmm mask vnni" "skylake-avx512|zmm|vnni" - "cascadelake|zmm|" + "cascadelake|zmm vnni|" ) set(svs_budget_found FALSE) diff --git a/cmake/dispatch-surface.cmake b/cmake/dispatch-surface.cmake index 971b1785..7cbe19b2 100644 --- a/cmake/dispatch-surface.cmake +++ b/cmake/dispatch-surface.cmake @@ -91,11 +91,29 @@ set(SVS_SUPPORTED_DIMS 64 96 100 128 160 200 512 768) # include/svs/multi-arch/x86/.cpp, and the object # library it is compiled into # -# Adding a level here also requires a `SVS_TYPE_PAIRS_` list in -# include/svs/multi-arch/x86/preprocessor.h, saying which element-type pairs -# that level has kernels for. That is deliberately not configured here: a type -# pair exists because an implementation exists for it. +# Unlike the extent list above, this list is not a knob. A level exists because +# kernels, a translation unit and a runtime check for it exist, so changing it +# means changing code: +# +# - a `SVS_TYPE_PAIRS_` list in +# include/svs/multi-arch/x86/preprocessor.h, saying which element-type pairs +# the level has kernels for +# - the specializations themselves, in the three distance headers +# - a branch in the entry points, and the CPUID check it tests +# +# The generated header defines `SVS_ISA_LEVEL_` for each level listed +# here, which is how the entry points tell a level that is present from one that +# is not; dropping AVX2 or AVX512 from this list is a compile error rather than a +# silent fall back to unvectorized code. +# +# On the instruction budgets in particular: `skylake-avx512` rather than +# `cascadelake` for the AVX512 level, because that level promises AVX-512 F/BW/DQ +# and nothing more. `cascadelake` also enables AVX512-VNNI, which the compiler is +# then free to emit into kernels that run on any AVX512F host -- a Skylake-SP +# among them, where `vpdpbusd` does not exist. VNNI is its own level instead, so +# that the budget and the promise line up. set(SVS_ISA_LEVELS "AVX2|haswell|avx2" - "AVX512|cascadelake|avx512" + "AVX512|skylake-avx512|avx512" + "AVX512_VNNI|cascadelake|vnni" ) diff --git a/cmake/generate-dispatch-surface.cmake b/cmake/generate-dispatch-surface.cmake index b1aa9a7c..75a06c69 100644 --- a/cmake/generate-dispatch-surface.cmake +++ b/cmake/generate-dispatch-surface.cmake @@ -68,6 +68,7 @@ string(APPEND SVS_GEN_DIM_LOOP " /* end */") set(SVS_GEN_TARGET_LOOP "\\\n") set(SVS_GEN_LEVEL_LOOP "\\\n") +set(SVS_GEN_LEVEL_DEFINES "") set(SVS_DISPATCH_TU_SPECS) foreach(level_spec IN LISTS SVS_ISA_LEVELS) string(REPLACE "|" ";" level_fields "${level_spec}") @@ -76,6 +77,8 @@ foreach(level_spec IN LISTS SVS_ISA_LEVELS) list(GET level_fields 2 infix) string(APPEND SVS_GEN_LEVEL_LOOP " M(${level}) \\\n") + string(APPEND SVS_GEN_LEVEL_DEFINES "#define SVS_ISA_LEVEL_${level} 1\n") + foreach(dim IN LISTS SVS_DIM_LIST) string(APPEND SVS_GEN_TARGET_LOOP " M(${dim}, ${level}) \\\n") endforeach() @@ -91,6 +94,7 @@ foreach(level_spec IN LISTS SVS_ISA_LEVELS) endforeach() string(APPEND SVS_GEN_TARGET_LOOP " /* end */") string(APPEND SVS_GEN_LEVEL_LOOP " /* end */") +string(STRIP "${SVS_GEN_LEVEL_DEFINES}" SVS_GEN_LEVEL_DEFINES) ##### ##### Emit the header diff --git a/cmake/multi-arch.cmake b/cmake/multi-arch.cmake index 74be6c57..e10e6e6c 100644 --- a/cmake/multi-arch.cmake +++ b/cmake/multi-arch.cmake @@ -30,7 +30,9 @@ foreach(tu_spec IN LISTS SVS_DISPATCH_TU_SPECS) add_library(${lib_name} INTERFACE) target_compile_options(${lib_name} INTERFACE -march=${arch} -mtune=${arch}) - set(obj_name ${arch}_obj) + # Named after the level, not the -march: more than one level can share an + # instruction budget, and it is the level that says what is inside. + set(obj_name ${infix}_obj) add_library(${obj_name} OBJECT ${src}) target_link_libraries( ${obj_name} PRIVATE ${SVS_LIB} svs::compile_options fmt::fmt ${lib_name} diff --git a/cmake/templates/dispatch_surface.h.in b/cmake/templates/dispatch_surface.h.in index 64acd976..ffbb9c50 100644 --- a/cmake/templates/dispatch_surface.h.in +++ b/cmake/templates/dispatch_surface.h.in @@ -37,4 +37,8 @@ // enumerators without a translation unit are absent: this is the surface. #define SVS_FOR_EACH_ISA_LEVEL(M) @SVS_GEN_LEVEL_LOOP@ +// One per ISA level in the surface, and only those: a level absent from it has no +// instantiations. Tested with `defined` so a -Wundef build stays quiet. +@SVS_GEN_LEVEL_DEFINES@ + // clang-format on diff --git a/include/svs/core/distance/cosine.h b/include/svs/core/distance/cosine.h index 4cd8e378..75f617ab 100644 --- a/include/svs/core/distance/cosine.h +++ b/include/svs/core/distance/cosine.h @@ -47,6 +47,18 @@ class CosineSimilarity { public: template static constexpr float compute(const Ea* a, const Eb* b, float a_norm, size_t N) { + if constexpr (has_vnni_kernel) { + if (__builtin_expect( + svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1 + )) { + return CosineSimilarityImpl< + Dynamic, + Ea, + Eb, + AVX_AVAILABILITY::AVX512_VNNI>:: + compute(a, b, a_norm, lib::MaybeStatic(N)); + } + } if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512f_supported(), 1)) { return CosineSimilarityImpl::compute( a, b, a_norm, lib::MaybeStatic(N) @@ -65,6 +77,23 @@ class CosineSimilarity { template static constexpr float compute(const Ea* a, const Eb* b, float a_norm) { + if constexpr (has_vnni_kernel) { + if (__builtin_expect( + svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1 + )) { + if constexpr (is_dim_supported()) { + return CosineSimilarityImpl:: + compute(a, b, a_norm, lib::MaybeStatic()); + } else { + return CosineSimilarityImpl< + Dynamic, + Ea, + Eb, + AVX_AVAILABILITY::AVX512_VNNI>:: + compute(a, b, a_norm, lib::MaybeStatic(N)); + } + } + } if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512f_supported(), 1)) { if constexpr (is_dim_supported()) { return CosineSimilarityImpl::compute( @@ -255,38 +284,70 @@ template <> struct CosineFloatOp<16> : public svs::simd::ConvertToFloat<16> { } }; -// Small Integers +// Small Integers, with VNNI SVS_VALIDATE_BOOL_ENV(SVS_AVX512_VNNI) #if SVS_AVX512_VNNI template -struct CosineSimilarityImpl { +struct CosineSimilarityImpl { SVS_NOINLINE static float compute(const int8_t* a, const int8_t* b, float a_norm, lib::MaybeStatic length) { - if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1)) { - auto sum = _mm512_setzero_epi32(); - auto bnorm_accum = _mm512_setzero_epi32(); - auto mask = create_mask<32>(length); - auto all = no_mask<32>(); - - for (size_t j = 0; j < length.size(); j += 32) { - auto temp_a = - _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, a + j); - auto va = _mm512_cvtepi8_epi16(temp_a); - - auto temp_b = - _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, b + j); - auto vb = _mm512_cvtepi8_epi16(temp_b); - - bnorm_accum = _mm512_dpwssd_epi32(bnorm_accum, vb, vb); - sum = _mm512_dpwssd_epi32(sum, va, vb); - } + auto sum = _mm512_setzero_epi32(); + auto bnorm_accum = _mm512_setzero_epi32(); + auto mask = create_mask<32>(length); + auto all = no_mask<32>(); + + for (size_t j = 0; j < length.size(); j += 32) { + auto temp_a = + _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, a + j); + auto va = _mm512_cvtepi8_epi16(temp_a); + + auto temp_b = + _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, b + j); + auto vb = _mm512_cvtepi8_epi16(temp_b); + + bnorm_accum = _mm512_dpwssd_epi32(bnorm_accum, vb, vb); + sum = _mm512_dpwssd_epi32(sum, va, vb); + } + + float b_norm = std::sqrt(static_cast(_mm512_reduce_add_epi32(bnorm_accum))); + return lib::narrow_cast(_mm512_reduce_add_epi32(sum)) / (a_norm * b_norm); + } +}; - float b_norm = - std::sqrt(static_cast(_mm512_reduce_add_epi32(bnorm_accum))); - return lib::narrow_cast(_mm512_reduce_add_epi32(sum)) / - (a_norm * b_norm); +template +struct CosineSimilarityImpl { + SVS_NOINLINE static float + compute(const uint8_t* a, const uint8_t* b, float a_norm, lib::MaybeStatic length) { + auto sum = _mm512_setzero_epi32(); + auto bnorm_accum = _mm512_setzero_epi32(); + auto mask = create_mask<32>(length); + auto all = no_mask<32>(); + + for (size_t j = 0; j < length.size(); j += 32) { + auto temp_a = + _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, a + j); + auto va = _mm512_cvtepu8_epi16(temp_a); + + auto temp_b = + _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, b + j); + auto vb = _mm512_cvtepu8_epi16(temp_b); + + bnorm_accum = _mm512_dpwssd_epi32(bnorm_accum, vb, vb); + sum = _mm512_dpwssd_epi32(sum, va, vb); } - // Fallback to AVX512 + float b_norm = std::sqrt(static_cast(_mm512_reduce_add_epi32(bnorm_accum))); + return lib::narrow_cast(_mm512_reduce_add_epi32(sum)) / (a_norm * b_norm); + } +}; + +#endif + +// Must stay outside the SVS_AVX512_VNNI guard: avx512.cpp compiles with that macro +// at 0, and hiding this specialization there would silently select the generic kernel. +template +struct CosineSimilarityImpl { + SVS_NOINLINE static float + compute(const int8_t* a, const int8_t* b, float a_norm, lib::MaybeStatic length) { auto [sum, norm] = simd::generic_simd_op(CosineFloatOp<16>(), a, b, length); return sum / (std::sqrt(norm) * a_norm); } @@ -296,37 +357,11 @@ template struct CosineSimilarityImpl { SVS_NOINLINE static float compute(const uint8_t* a, const uint8_t* b, float a_norm, lib::MaybeStatic length) { - if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1)) { - auto sum = _mm512_setzero_epi32(); - auto bnorm_accum = _mm512_setzero_epi32(); - auto mask = create_mask<32>(length); - auto all = no_mask<32>(); - - for (size_t j = 0; j < length.size(); j += 32) { - auto temp_a = - _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, a + j); - auto va = _mm512_cvtepu8_epi16(temp_a); - - auto temp_b = - _mm256_maskz_loadu_epi8(islast<32>(length, j) ? mask : all, b + j); - auto vb = _mm512_cvtepu8_epi16(temp_b); - - bnorm_accum = _mm512_dpwssd_epi32(bnorm_accum, vb, vb); - sum = _mm512_dpwssd_epi32(sum, va, vb); - } - float b_norm = - std::sqrt(static_cast(_mm512_reduce_add_epi32(bnorm_accum))); - return lib::narrow_cast(_mm512_reduce_add_epi32(sum)) / - (a_norm * b_norm); - } - // Fallback to AVX512 auto [sum, norm] = simd::generic_simd_op(CosineFloatOp<16>(), a, b, length); return sum / (std::sqrt(norm) * a_norm); } }; -#endif - // Floating and Mixed Types template struct CosineSimilarityImpl { SVS_NOINLINE static float diff --git a/include/svs/core/distance/dispatch_surface.h b/include/svs/core/distance/dispatch_surface.h index a357ac0a..152e6d92 100644 --- a/include/svs/core/distance/dispatch_surface.h +++ b/include/svs/core/distance/dispatch_surface.h @@ -60,6 +60,15 @@ M(512, AVX512) \ M(768, AVX512) \ M(svs::Dynamic, AVX512) \ + M(64, AVX512_VNNI) \ + M(96, AVX512_VNNI) \ + M(100, AVX512_VNNI) \ + M(128, AVX512_VNNI) \ + M(160, AVX512_VNNI) \ + M(200, AVX512_VNNI) \ + M(512, AVX512_VNNI) \ + M(768, AVX512_VNNI) \ + M(svs::Dynamic, AVX512_VNNI) \ /* end */ // Invokes M(isa_level) once per ISA level, weakest first. AVX_AVAILABILITY @@ -67,6 +76,13 @@ #define SVS_FOR_EACH_ISA_LEVEL(M) \ M(AVX2) \ M(AVX512) \ + M(AVX512_VNNI) \ /* end */ +// One per ISA level in the surface, and only those: a level absent from it has no +// instantiations. Tested with `defined` so a -Wundef build stays quiet. +#define SVS_ISA_LEVEL_AVX2 1 +#define SVS_ISA_LEVEL_AVX512 1 +#define SVS_ISA_LEVEL_AVX512_VNNI 1 + // clang-format on diff --git a/include/svs/core/distance/distance_core.h b/include/svs/core/distance/distance_core.h index 5d50dee1..80ba089b 100644 --- a/include/svs/core/distance/distance_core.h +++ b/include/svs/core/distance/distance_core.h @@ -24,12 +24,41 @@ // The extent list and the ISA levels, generated from cmake/dispatch-surface.cmake. #include "svs/core/distance/dispatch_surface.h" +// Needed here and not only where the kernels are declared: the entry points must +// not dispatch to a level with no kernel for the pair in hand. +#if defined(__x86_64__) +#include "svs/multi-arch/x86/preprocessor.h" + +// Dispatched to for every type pair, so a surface without them would leave each +// consumer to instantiate the kernels itself, from the generic template. +#if !defined(SVS_ISA_LEVEL_AVX2) || !defined(SVS_ISA_LEVEL_AVX512) +#error "the x86 dispatch surface must declare the AVX2 and AVX512 ISA levels" +#endif +#endif + #include #include namespace svs::distance { -enum class AVX_AVAILABILITY { NONE, AVX2, AVX512 }; +/// The runtime ISA levels the library compiles distance kernels for. +/// +/// Each is a promise about the host, checked once in the entry point; the kernels +/// branch on nothing. Append new levels -- renumbering changes mangled names. +enum class AVX_AVAILABILITY { NONE, AVX2, AVX512, AVX512_VNNI }; + +/// Whether (Ea, Eb) has a kernel at AVX_AVAILABILITY::AVX512_VNNI. +/// +/// False where there is no such kernel -- a float-promoting pair, or any pair when +/// the surface omits the level. Dispatching anyway instantiates the generic template. +template inline constexpr bool has_vnni_kernel = false; + +#if defined(__x86_64__) && defined(SVS_ISA_LEVEL_AVX512_VNNI) +#define SVS_MARK_VNNI_PAIR(Ea, Eb, ...) \ + template <> inline constexpr bool has_vnni_kernel = true; +SVS_TYPE_PAIRS_AVX512_VNNI(SVS_MARK_VNNI_PAIR, ) +#undef SVS_MARK_VNNI_PAIR +#endif /// The extents that have a fixed-extent kernel, including svs::Dynamic. #define SVS_DIM_LIST_ENTRY(N) N, diff --git a/include/svs/core/distance/euclidean.h b/include/svs/core/distance/euclidean.h index a2fa6848..1d373419 100644 --- a/include/svs/core/distance/euclidean.h +++ b/include/svs/core/distance/euclidean.h @@ -86,6 +86,15 @@ class L2 { public: template static constexpr float compute(const Ea* a, const Eb* b, size_t N) { + if constexpr (has_vnni_kernel) { + if (__builtin_expect( + svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1 + )) { + return L2Impl::compute( + a, b, lib::MaybeStatic(N) + ); + } + } if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512f_supported(), 1)) { return L2Impl::compute( a, b, lib::MaybeStatic(N) @@ -103,6 +112,21 @@ class L2 { template static constexpr float compute(const Ea* a, const Eb* b) { + if constexpr (has_vnni_kernel) { + if (__builtin_expect( + svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1 + )) { + if constexpr (is_dim_supported()) { + return L2Impl::compute( + a, b, lib::MaybeStatic() + ); + } else { + return L2Impl::compute( + a, b, lib::MaybeStatic(N) + ); + } + } + } if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512f_supported(), 1)) { if constexpr (is_dim_supported()) { return L2Impl::compute( @@ -258,7 +282,7 @@ template <> struct L2FloatOp<16> : public svs::simd::ConvertToFloat<16> { static float reduce(__m512 x) { return _mm512_reduce_add_ps(x); } }; -// Small Integers +// Small Integers, with VNNI SVS_VALIDATE_BOOL_ENV(SVS_AVX512_VNNI) #if SVS_AVX512_VNNI @@ -289,14 +313,27 @@ template <> struct L2VNNIOp : public svs::simd::ConvertForVNNI struct L2Impl { + SVS_NOINLINE static float + compute(const int8_t* a, const int8_t* b, lib::MaybeStatic length) { + return simd::generic_simd_op(L2VNNIOp(), a, b, length); + } +}; + +template struct L2Impl { + SVS_NOINLINE static float + compute(const uint8_t* a, const uint8_t* b, lib::MaybeStatic length) { + return simd::generic_simd_op(L2VNNIOp(), a, b, length); + } +}; + +#endif + +// Must stay outside the SVS_AVX512_VNNI guard: avx512.cpp compiles with that macro +// at 0, and hiding this specialization there would silently select the generic kernel. template struct L2Impl { SVS_NOINLINE static float compute(const int8_t* a, const int8_t* b, lib::MaybeStatic length) { - if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1)) { - return simd::generic_simd_op(L2VNNIOp(), a, b, length); - } - // fallback to AVX512 return simd::generic_simd_op(L2FloatOp<16>{}, a, b, length); } }; @@ -304,16 +341,10 @@ template struct L2Impl { template struct L2Impl { SVS_NOINLINE static float compute(const uint8_t* a, const uint8_t* b, lib::MaybeStatic length) { - if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1)) { - return simd::generic_simd_op(L2VNNIOp(), a, b, length); - } - // fallback to AVX512 return simd::generic_simd_op(L2FloatOp<16>{}, a, b, length); } }; -#endif - // Floating and Mixed Types template struct L2Impl { SVS_NOINLINE static float diff --git a/include/svs/core/distance/inner_product.h b/include/svs/core/distance/inner_product.h index 14293cb3..c7d958a0 100644 --- a/include/svs/core/distance/inner_product.h +++ b/include/svs/core/distance/inner_product.h @@ -46,6 +46,15 @@ class IP { public: template static constexpr float compute(const Ea* a, const Eb* b, size_t N) { + if constexpr (has_vnni_kernel) { + if (__builtin_expect( + svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1 + )) { + return IPImpl::compute( + a, b, lib::MaybeStatic(N) + ); + } + } if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512f_supported(), 1)) { return IPImpl::compute( a, b, lib::MaybeStatic(N) @@ -63,6 +72,21 @@ class IP { template static constexpr float compute(const Ea* a, const Eb* b) { + if constexpr (has_vnni_kernel) { + if (__builtin_expect( + svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1 + )) { + if constexpr (is_dim_supported()) { + return IPImpl::compute( + a, b, lib::MaybeStatic() + ); + } else { + return IPImpl::compute( + a, b, lib::MaybeStatic(N) + ); + } + } + } if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512f_supported(), 1)) { if constexpr (is_dim_supported()) { return IPImpl::compute( @@ -215,7 +239,7 @@ template <> struct IPFloatOp<16> : public svs::simd::ConvertToFloat<16> { static float reduce(__m512 x) { return _mm512_reduce_add_ps(x); } }; -// Small Integers +// Small Integers, with VNNI SVS_VALIDATE_BOOL_ENV(SVS_AVX512_VNNI) #if SVS_AVX512_VNNI @@ -243,14 +267,27 @@ template <> struct IPVNNIOp : public svs::simd::ConvertForVNNI struct IPImpl { + SVS_NOINLINE static float + compute(const int8_t* a, const int8_t* b, lib::MaybeStatic length) { + return simd::generic_simd_op(IPVNNIOp(), a, b, length); + } +}; + +template struct IPImpl { + SVS_NOINLINE static float + compute(const uint8_t* a, const uint8_t* b, lib::MaybeStatic length) { + return simd::generic_simd_op(IPVNNIOp(), a, b, length); + } +}; + +#endif + +// Must stay outside the SVS_AVX512_VNNI guard: avx512.cpp compiles with that macro +// at 0, and hiding this specialization there would silently select the generic kernel. template struct IPImpl { SVS_NOINLINE static float compute(const int8_t* a, const int8_t* b, lib::MaybeStatic length) { - if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1)) { - return simd::generic_simd_op(IPVNNIOp(), a, b, length); - } - // fallback to AVX512 return svs::simd::generic_simd_op(IPFloatOp<16>{}, a, b, length); } }; @@ -258,16 +295,10 @@ template struct IPImpl { template struct IPImpl { SVS_NOINLINE static float compute(const uint8_t* a, const uint8_t* b, lib::MaybeStatic length) { - if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1)) { - return simd::generic_simd_op(IPVNNIOp(), a, b, length); - } - // fallback to AVX512 return svs::simd::generic_simd_op(IPFloatOp<16>{}, a, b, length); } }; -#endif - // Floating and Mixed Types template struct IPImpl { SVS_NOINLINE static float diff --git a/include/svs/multi-arch/x86/preprocessor.h b/include/svs/multi-arch/x86/preprocessor.h index 75941b70..82eccd84 100644 --- a/include/svs/multi-arch/x86/preprocessor.h +++ b/include/svs/multi-arch/x86/preprocessor.h @@ -55,6 +55,12 @@ #define SVS_TYPE_PAIRS_AVX2 SVS_FOR_EACH_TYPE_PAIR #define SVS_TYPE_PAIRS_AVX512 SVS_FOR_EACH_TYPE_PAIR +// Only these two: every other pair promotes to float, where VNNI has nothing to +// offer. svs::distance::has_vnni_kernel is generated from this same list. +#define SVS_TYPE_PAIRS_AVX512_VNNI(M, ...) \ + M(int8_t, int8_t, __VA_ARGS__) \ + M(uint8_t, uint8_t, __VA_ARGS__) + ///// ///// Instantiation. ///// diff --git a/include/svs/multi-arch/x86/vnni.cpp b/include/svs/multi-arch/x86/vnni.cpp new file mode 100644 index 00000000..9c2810b0 --- /dev/null +++ b/include/svs/multi-arch/x86/vnni.cpp @@ -0,0 +1,32 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#if defined(__x86_64__) +#include "svs/core/distance/cosine.h" +#include "svs/core/distance/euclidean.h" +#include "svs/core/distance/inner_product.h" + +namespace svs::distance { + +// Define every kernel for this ISA level at every generated extent. Extents come +// from cmake/dispatch-surface.cmake, type pairs from multi-arch/x86/preprocessor.h. +#define SVS_DEFINE_FOR_DIM(DIM) SVS_INSTANTIATE_DISTANCES(template, DIM, AVX512_VNNI) +SVS_FOR_EACH_SUPPORTED_DIM(SVS_DEFINE_FOR_DIM) +#undef SVS_DEFINE_FOR_DIM + +} // namespace svs::distance + +#endif diff --git a/tests/cmake/dispatch-surface/valid-reduced.cmake b/tests/cmake/dispatch-surface/valid-reduced.cmake index 4a9bfcc0..dcabc8e5 100644 --- a/tests/cmake/dispatch-surface/valid-reduced.cmake +++ b/tests/cmake/dispatch-surface/valid-reduced.cmake @@ -13,13 +13,18 @@ # limitations under the License. # A surface that shares no extent with the default declaration, so a build using -# it cannot accidentally pass by reusing a committed header. Both ISA levels are -# kept so that runtime dispatch is still exercised on an AVX-512 host. +# it cannot accidentally pass by reusing a committed header. +# +# The extent list is what this fixture varies. The ISA levels are copied from the +# default declaration verbatim, instruction budgets included: they are not +# configuration, and a build with a different level set would be testing a library +# nobody ships. # # Built and tested by the `non-default surface` CI job. set(SVS_SUPPORTED_DIMS 32 384) set(SVS_ISA_LEVELS "AVX2|haswell|avx2" - "AVX512|cascadelake|avx512" + "AVX512|skylake-avx512|avx512" + "AVX512_VNNI|cascadelake|vnni" ) diff --git a/tests/multi-arch/CMakeLists.txt b/tests/multi-arch/CMakeLists.txt index d52c1bc2..744dfc40 100644 --- a/tests/multi-arch/CMakeLists.txt +++ b/tests/multi-arch/CMakeLists.txt @@ -112,7 +112,7 @@ if(svs_objdump) NAME dispatch_instructions_${svs_infix} COMMAND "${CMAKE_COMMAND}" - "-DSVS_OBJECT=$" + "-DSVS_OBJECT=$" "-DSVS_LEVEL=${svs_level}" "-DSVS_ARCH=${svs_arch}" "-DSVS_OBJDUMP=${svs_objdump}" diff --git a/tests/multi-arch/x86/entry_probe.cpp b/tests/multi-arch/x86/entry_probe.cpp index 567ae678..f5571910 100644 --- a/tests/multi-arch/x86/entry_probe.cpp +++ b/tests/multi-arch/x86/entry_probe.cpp @@ -29,6 +29,7 @@ #include #include #include +#include #include #include @@ -98,9 +99,11 @@ int main(int argc, char** argv) { // breaks on this extent's kernels to see which level the call enters. std::printf("expect-level %d\n", svs_test::expected_level()); std::printf("probe-extent %zu\n", entry_report_dim); + // int8/int8, mangled `aa` and matched in cmake/check-dispatch-execution.cmake: + // a float-promoting pair has no AVX512_VNNI kernel and routes a level lower. std::printf( "one-call %f\n", - static_cast(entry_one()) + static_cast(entry_one()) ); return 0; } diff --git a/tests/multi-arch/x86/host_levels.h b/tests/multi-arch/x86/host_levels.h index a94c4297..b8372655 100644 --- a/tests/multi-arch/x86/host_levels.h +++ b/tests/multi-arch/x86/host_levels.h @@ -38,6 +38,10 @@ template <> inline bool host_satisfies() { return svs::detail::avx_runtime_flags.is_avx512f_supported(); } +template <> inline bool host_satisfies() { + return svs::detail::avx_runtime_flags.is_avx512vnni_supported(); +} + /// The level the distance entry points must choose on this host. /// /// The entry points test the strongest level first, so their choice is the diff --git a/tests/multi-arch/x86/link_probe.cpp b/tests/multi-arch/x86/link_probe.cpp index 443c87e6..ea540a92 100644 --- a/tests/multi-arch/x86/link_probe.cpp +++ b/tests/multi-arch/x86/link_probe.cpp @@ -86,9 +86,11 @@ template svs::lib::MaybeStatic probe_length() { buffer(), buffer(), 1.0F, probe_length() \ ); -#define SVS_PROBE_TARGET(N, LEVEL) \ - if (host_satisfies()) { \ - SVS_FOR_EACH_TYPE_PAIR(SVS_PROBE_ONE, N, LEVEL) \ +// The level's own type-pair list, not the full one: naming a pair the level has no +// kernel for would instantiate the generic template right here. +#define SVS_PROBE_TARGET(N, LEVEL) \ + if (host_satisfies()) { \ + SVS_TYPE_PAIRS_FOR(LEVEL, SVS_PROBE_ONE, N, LEVEL) \ } float probe_all() {