fix(dispatch): Make VNNI its own ISA level and lower AVX512 to skylake-avx512

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()
 

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)

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/<infix>.cpp, and the object
 #                      library it is compiled into
 #
-# Adding a level here also requires a `SVS_TYPE_PAIRS_<enumerator>` 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_<enumerator>` 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_<enumerator>` 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"
 )

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

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}

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

include/svs/core/distance/cosine.h

@@ -47,6 +47,18 @@ class CosineSimilarity {
   public:
     template <typename Ea, typename Eb>
     static constexpr float compute(const Ea* a, const Eb* b, float a_norm, size_t N) {
+        if constexpr (has_vnni_kernel<Ea, Eb>) {
+            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<Dynamic, Ea, Eb, AVX_AVAILABILITY::AVX512>::compute(
                 a, b, a_norm, lib::MaybeStatic(N)
@@ -65,6 +77,23 @@ class CosineSimilarity {
 
     template <size_t N, typename Ea, typename Eb>
     static constexpr float compute(const Ea* a, const Eb* b, float a_norm) {
+        if constexpr (has_vnni_kernel<Ea, Eb>) {
+            if (__builtin_expect(
+                    svs::detail::avx_runtime_flags.is_avx512vnni_supported(), 1
+                )) {
+                if constexpr (is_dim_supported<N>()) {
+                    return CosineSimilarityImpl<N, Ea, Eb, AVX_AVAILABILITY::AVX512_VNNI>::
+                        compute(a, b, a_norm, lib::MaybeStatic<N>());
+                } 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<N>()) {
                 return CosineSimilarityImpl<N, Ea, Eb, AVX_AVAILABILITY::AVX512>::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 <size_t N>
-struct CosineSimilarityImpl<N, int8_t, int8_t, AVX_AVAILABILITY::AVX512> {
+struct CosineSimilarityImpl<N, int8_t, int8_t, AVX_AVAILABILITY::AVX512_VNNI> {
     SVS_NOINLINE static float
     compute(const int8_t* a, const int8_t* b, float a_norm, lib::MaybeStatic<N> 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<float>(_mm512_reduce_add_epi32(bnorm_accum)));
+        return lib::narrow_cast<float>(_mm512_reduce_add_epi32(sum)) / (a_norm * b_norm);
+    }
+};
 
-            float b_norm =
-                std::sqrt(static_cast<float>(_mm512_reduce_add_epi32(bnorm_accum)));
-            return lib::narrow_cast<float>(_mm512_reduce_add_epi32(sum)) /
-                   (a_norm * b_norm);
+template <size_t N>
+struct CosineSimilarityImpl<N, uint8_t, uint8_t, AVX_AVAILABILITY::AVX512_VNNI> {
+    SVS_NOINLINE static float
+    compute(const uint8_t* a, const uint8_t* b, float a_norm, lib::MaybeStatic<N> 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<float>(_mm512_reduce_add_epi32(bnorm_accum)));
+        return lib::narrow_cast<float>(_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 <size_t N>
+struct CosineSimilarityImpl<N, int8_t, int8_t, AVX_AVAILABILITY::AVX512> {
+    SVS_NOINLINE static float
+    compute(const int8_t* a, const int8_t* b, float a_norm, lib::MaybeStatic<N> 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 <size_t N>
 struct CosineSimilarityImpl<N, uint8_t, uint8_t, AVX_AVAILABILITY::AVX512> {
     SVS_NOINLINE static float
     compute(const uint8_t* a, const uint8_t* b, float a_norm, lib::MaybeStatic<N> 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<float>(_mm512_reduce_add_epi32(bnorm_accum)));
-            return lib::narrow_cast<float>(_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 <size_t N> struct CosineSimilarityImpl<N, float, float, AVX_AVAILABILITY::AVX512> {
     SVS_NOINLINE static float

include/svs/core/distance/dispatch_surface.h

@@ -60,13 +60,29 @@
     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
 // enumerators without a translation unit are absent: this is the surface.
 #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

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 <cmath>
 #include <span>
 
 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 <typename Ea, typename Eb> 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<Ea, Eb> = 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,