tests: cover the AVX2/AVX-512 dot-product kernels

tests/CMakeLists.txt

@@ -75,3 +75,49 @@ include(GoogleTest)
 # every test case can exceed the 5s default, causing a spurious build failure.
 # Bump the timeout well above that to keep discovery robust.
 gtest_discover_tests(nsparse_test DISCOVERY_TIMEOUT 120)
+
+# SIMD kernel equivalence tests. distance_test.cpp above runs through the
+# flag-less `nsparse`, so it only ever exercises the scalar distance.h; the SIMD
+# kernels that ship in production (nsparse_avx2/avx512, selected at load time by
+# the Python bindings) are otherwise untested. Each binary below compiles
+# distance_kernel_equivalence_test.cpp with one ISA's -m flags so distance_simd.h
+# resolves to that SIMD kernel, and checks it against a plain reference.
+#
+# Because the whole binary is compiled for the ISA, its static initializers can
+# execute that ISA before main() -- so it must not even be launched on a CPU
+# that lacks the ISA (that would SIGILL, ahead of any in-process guard). ctest
+# therefore invokes each binary through run_if_isa_supported.cmake, which checks
+# the host CPU first and reports a skip otherwise; this also keeps the AVX-512
+# binary from faulting on an AVX-512-less host.
+#
+# Each binary links the matching per-ISA library (nsparse_avx2/avx512), so every
+# translation unit -- the test and the library -- is compiled for the same ISA:
+# the shared nsparse::detail kernel symbols then have one consistent definition
+# (no ODR violation), and SparseVectors' out-of-line symbols resolve. x86 only
+# for now; the NEON/SVE kernels want an ARM CI host and are a follow-up.
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64|AMD64|i[3-6]86)" AND NOT WIN32)
+    # isa: label (also NSPARSE_EXPECT_ISA); cpu_flag: the /proc/cpuinfo token to
+    # gate the launch on; ARGN: the -m compile flags.
+    function(nsparse_add_kernel_test isa cpu_flag)
+        set(_tgt nsparse_${isa}_kernel_test)
+        add_executable(${_tgt} distance_kernel_equivalence_test.cpp)
+        target_include_directories(${_tgt} PRIVATE ${PROJECT_SOURCE_DIR})
+        target_compile_options(${_tgt} PRIVATE ${ARGN})
+        target_compile_definitions(${_tgt} PRIVATE NSPARSE_EXPECT_ISA="${isa}")
+        target_link_libraries(${_tgt} PRIVATE
+            nsparse_${isa}
+            OpenMP::OpenMP_CXX
+            GTest::gtest_main
+            absl::flat_hash_set
+        )
+        add_test(NAME ${_tgt}
+                 COMMAND ${CMAKE_COMMAND}
+                         -DTEST_EXE=$<TARGET_FILE:${_tgt}>
+                         -DCPU_FLAG=${cpu_flag}
+                         -P ${CMAKE_CURRENT_SOURCE_DIR}/run_if_isa_supported.cmake)
+    endfunction()
+
+    nsparse_add_kernel_test(avx2 avx2 -mavx2 -mfma -mf16c -mpopcnt)
+    nsparse_add_kernel_test(avx512 avx512f -mavx2 -mfma -mf16c -mavx512f
+                            -mavx512cd -mavx512vl -mavx512dq -mavx512bw -mpopcnt)
+endif()

tests/distance_kernel_equivalence_test.cpp

@@ -0,0 +1,200 @@
+/**
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+// The dot-product kernels are selected at compile time by distance_simd.h
+// (`#if __AVX2__ / __AVX512F__ / ...`), and the main test binary links the
+// flag-less library, so distance_test.cpp only ever exercises the scalar
+// distance.h. This binary is compiled once per ISA with the matching -m flags
+// (see tests/CMakeLists.txt), so `distance_simd.h` here resolves to the SIMD
+// kernel, and each dot product is checked against a plain, independent
+// reference. Scalar and SIMD kernels share the same nsparse::detail symbol
+// names, so they cannot be linked into one binary -- hence the reference is
+// computed locally here rather than pulled from distance.h.
+//
+// ctest gates the launch on CPU support (run_if_isa_supported.cmake), because an
+// ISA-specialized binary can execute the ISA in a static initializer before any
+// in-process check runs -- so gating has to happen before it starts, not inside
+// it. It also fails to build if CMake's requested ISA and the ISA the
+// preprocessor selected disagree (a missing -m flag), so it can never silently
+// degrade to scalar.
+
+#include <gtest/gtest.h>
+
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <random>
+#include <set>
+#include <string_view>
+#include <vector>
+
+#include "nsparse/types.h"
+#include "nsparse/utils/distance_simd.h"
+
+#ifndef NSPARSE_EXPECT_ISA
+#define NSPARSE_EXPECT_ISA "scalar"
+#endif
+
+namespace nsparse {
+namespace {
+
+// The kernel family this translation unit was compiled for, decided by the same
+// macros distance_simd.h dispatches on.
+#if defined(__AVX512F__)
+constexpr std::string_view kIsa = "avx512";
+#elif defined(__AVX2__)
+constexpr std::string_view kIsa = "avx2";
+#else
+constexpr std::string_view kIsa = "scalar";
+#endif
+
+// A missing -m flag would leave the preprocessor on the scalar branch while
+// CMake believed it built a SIMD binary; catch that at build time.
+static_assert(kIsa == std::string_view(NSPARSE_EXPECT_ISA),
+              "compiled kernel ISA does not match the one CMake requested -- a "
+              "-m compile flag is missing");
+
+constexpr int kDim = 1024;
+
+// Independent references (not distance.h -- see the file header). Float sums in
+// double so the tolerance only has to absorb the kernel's own reordering;
+// integer sums are exact.
+double ref_float(const term_t* indices, const float* weights, size_t len,
+                 const float* dense) {
+    double sum = 0.0;
+    for (size_t i = 0; i < len; ++i) {
+        sum += static_cast<double>(weights[i]) *
+               static_cast<double>(dense[indices[i]]);
+    }
+    return sum;
+}
+
+template <class T>
+int64_t ref_int(const term_t* indices, const T* weights, size_t len,
+                const T* dense) {
+    int64_t sum = 0;
+    for (size_t i = 0; i < len; ++i) {
+        sum += static_cast<int64_t>(weights[i]) *
+               static_cast<int64_t>(dense[indices[i]]);
+    }
+    return sum;
+}
+
+// Distinct term ids in [0, kDim) with per-length weights and a dense buffer.
+// Integer magnitudes are kept modest so the int64 reference sum stays below
+// 2^24 and its float cast is exact -- any kernel discrepancy then shows up
+// exactly, rather than being hidden by float rounding.
+struct Inputs {
+    std::vector<term_t> indices;
+    std::vector<float> fweights, fdense;
+    std::vector<uint16_t> u16weights, u16dense;
+    std::vector<uint8_t> u8weights, u8dense;
+};
+
+Inputs make_inputs(size_t len, uint32_t seed) {
+    std::mt19937 rng(seed);
+    std::uniform_int_distribution<int> idx_dist(0, kDim - 1);
+    std::uniform_real_distribution<float> f_dist(-4.0F, 4.0F);
+    std::uniform_int_distribution<int> small_dist(0, 127);
+
+    Inputs in;
+    std::set<int> chosen;  // distinct, like a document's component ids
+    while (chosen.size() < len) {
+        chosen.insert(idx_dist(rng));
+    }
+    for (const int t : chosen) {
+        in.indices.push_back(static_cast<term_t>(t));
+    }
+    for (size_t i = 0; i < len; ++i) {
+        in.fweights.push_back(f_dist(rng));
+        in.u16weights.push_back(static_cast<uint16_t>(small_dist(rng)));
+        in.u8weights.push_back(static_cast<uint8_t>(small_dist(rng)));
+    }
+    in.fdense.resize(kDim);
+    in.u16dense.resize(kDim);
+    in.u8dense.resize(kDim);
+    for (int i = 0; i < kDim; ++i) {
+        in.fdense[i] = f_dist(rng);
+        in.u16dense[i] = static_cast<uint16_t>(small_dist(rng));
+        in.u8dense[i] = static_cast<uint8_t>(small_dist(rng));
+    }
+    return in;
+}
+
+// Lengths straddle the 8- and 16-wide SIMD strides (and AVX-512's 16-wide float
+// stride), so the scalar tail of every kernel is exercised.
+const std::vector<size_t> kLengths = {0,  1,  2,  3,  7,   8,   9,   15,
+                                      16, 17, 31, 32, 33,  63,  64,  65,
+                                      127, 128, 129, 255, 256, 257, 500};
+
+// ctest launches this binary only on a CPU that supports the compiled ISA
+// (run_if_isa_supported.cmake), so the kernels below always run when reached.
+class KernelEquivalence : public ::testing::Test {};
+
+TEST_F(KernelEquivalence, FloatDenseMatchesReference) {
+    for (const size_t len : kLengths) {
+        const Inputs in = make_inputs(len, /*seed=*/1000 + len);
+        const float got = detail::dot_product_float_dense(
+            in.indices.data(), in.fweights.data(), len, in.fdense.data());
+        const auto want = static_cast<float>(
+            ref_float(in.indices.data(), in.fweights.data(), len,
+                      in.fdense.data()));
+        const float tol = 1e-4F * (std::abs(want) + 1.0F);
+        EXPECT_NEAR(got, want, tol) << kIsa << " float, len=" << len;
+    }
+}
+
+TEST_F(KernelEquivalence, Uint16DenseMatchesReference) {
+    for (const size_t len : kLengths) {
+        const Inputs in = make_inputs(len, /*seed=*/2000 + len);
+        const float got = detail::dot_product_uint16_dense(
+            in.indices.data(), in.u16weights.data(), len, in.u16dense.data());
+        const auto want = static_cast<float>(
+            ref_int<uint16_t>(in.indices.data(), in.u16weights.data(), len,
+                              in.u16dense.data()));
+        EXPECT_FLOAT_EQ(got, want) << kIsa << " uint16, len=" << len;
+    }
+}
+
+TEST_F(KernelEquivalence, Uint8DenseMatchesReference) {
+    for (const size_t len : kLengths) {
+        const Inputs in = make_inputs(len, /*seed=*/3000 + len);
+        const float got = detail::dot_product_uint8_dense(
+            in.indices.data(), in.u8weights.data(), len, in.u8dense.data());
+        const auto want = static_cast<float>(
+            ref_int<uint8_t>(in.indices.data(), in.u8weights.data(), len,
+                             in.u8dense.data()));
+        EXPECT_FLOAT_EQ(got, want) << kIsa << " uint8, len=" << len;
+    }
+}
+
+// Full-width integer values (the overflow-prone path): the 8-/16-bit codes are
+// widened before multiply, so max-valued inputs must not wrap. Kept short so
+// the exact reference sum still lands below 2^24.
+TEST_F(KernelEquivalence, IntegerMaxValuesDoNotOverflow) {
+    constexpr size_t len = 16;
+    std::vector<term_t> indices(len);
+    for (size_t i = 0; i < len; ++i) {
+        indices[i] = static_cast<term_t>(i);
+    }
+    std::vector<uint8_t> u8w(len, 255), u8d(kDim, 255);
+    const float u8_got = detail::dot_product_uint8_dense(
+        indices.data(), u8w.data(), len, u8d.data());
+    EXPECT_FLOAT_EQ(u8_got, static_cast<float>(len) * 255.0F * 255.0F)
+        << kIsa << " uint8 max";
+
+    std::vector<uint16_t> u16w(len, 1000), u16d(kDim, 1000);
+    const float u16_got = detail::dot_product_uint16_dense(
+        indices.data(), u16w.data(), len, u16d.data());
+    EXPECT_FLOAT_EQ(u16_got, static_cast<float>(len) * 1000.0F * 1000.0F)
+        << kIsa << " uint16 max";
+}
+
+}  // namespace
+}  // namespace nsparse

tests/run_if_isa_supported.cmake

@@ -0,0 +1,46 @@
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+
+# Runs ${TEST_EXE} only when the host CPU exposes ${CPU_FLAG} (the /proc/cpuinfo
+# flag name, e.g. "avx2" or "avx512f"); otherwise reports a skip and succeeds.
+#
+# An ISA-specialized binary is compiled entirely with that ISA's -m flags, so
+# its static initializers (e.g. a namespace-scope object's constructor) can
+# execute the ISA before main() -- and thus before any in-process CPU check.
+# On a CPU without the ISA that is an illegal instruction (SIGILL). The only
+# safe guard is to decide whether to launch the binary at all, here, before it
+# starts. ctest invokes this via `cmake -P`.
+
+if(NOT DEFINED TEST_EXE OR NOT DEFINED CPU_FLAG)
+    message(FATAL_ERROR "run_if_isa_supported: TEST_EXE and CPU_FLAG required")
+endif()
+
+set(_supported FALSE)
+if(EXISTS "/proc/cpuinfo")
+    file(READ "/proc/cpuinfo" _cpuinfo)
+    # Match CPU_FLAG as a whole token (flags are separated by spaces, and the
+    # last one on a line by a newline -- both are non-word characters).
+    if(" ${_cpuinfo} " MATCHES "[^a-zA-Z0-9_]${CPU_FLAG}[^a-zA-Z0-9_]")
+        set(_supported TRUE)
+    endif()
+elseif(APPLE)
+    execute_process(COMMAND sysctl -a OUTPUT_VARIABLE _sysctl ERROR_QUIET)
+    string(TOLOWER "${_sysctl}" _sysctl)
+    if(_sysctl MATCHES "${CPU_FLAG}")
+        set(_supported TRUE)
+    endif()
+endif()
+
+if(NOT _supported)
+    message(STATUS "SKIP ${TEST_EXE}: host CPU does not expose '${CPU_FLAG}'")
+    return()
+endif()
+
+execute_process(COMMAND "${TEST_EXE}" RESULT_VARIABLE _rc)
+if(NOT _rc EQUAL 0)
+    message(FATAL_ERROR "${TEST_EXE} exited with ${_rc}")
+endif()