Add native SIMD testing and benchmark module

.github/workflows/unit-tests.yaml

@@ -45,12 +45,38 @@ jobs:
       - name: Set up GCC
         run: |
           sudo apt install -y gcc
-      - name: Install Meson and Ninja
+      - name: Install Meson, Ninja, and GTest
         run: |
-          sudo apt update && sudo apt install -y meson ninja-build
+          sudo apt update && sudo apt install -y meson ninja-build pkg-config libgtest-dev
       - uses: actions/checkout@v4
       - name: Initialize Git Submodules
         run: git submodule update --init
+
+      - name: Build test_simd_kernels (native C++)
+        working-directory: jvector-native/src/main/native
+        run: |
+          meson setup build --wipe
+          ninja -C build test_simd_kernels
+
+      - name: Run test_simd_kernels — no ISA cap (auto-detect)
+        if: matrix.max_isa == 'avx512f'
+        working-directory: jvector-native/src/main/native
+        run: ./build/test_simd_kernels
+
+      - name: Run test_simd_kernels — capped at avx2
+        if: matrix.max_isa == 'avx2'
+        working-directory: jvector-native/src/main/native
+        env:
+          JVECTOR_MAX_ISA: avx2
+        run: ./build/test_simd_kernels
+
+      - name: Run test_simd_kernels — capped at sse42
+        if: matrix.max_isa == 'sse42'
+        working-directory: jvector-native/src/main/native
+        env:
+          JVECTOR_MAX_ISA: sse42
+        run: ./build/test_simd_kernels
+
       - name: Set up JDK ${{ matrix.jdk }}
         uses: actions/setup-java@v3
         with:

jvector-native/pom.xml

@@ -141,7 +141,7 @@
                                         <argument>${native.buildtype}</argument>
                                     </arguments>
                                     <skip>false</skip>
-                                    <workingDirectory>${project.basedir}/src/main/native/</workingDirectory>
+                                    <workingDirectory>${project.basedir}/src/main/native/src/</workingDirectory>
                                 </configuration>
                             </execution>
                         </executions>

jvector-native/src/main/native/benchmarks/bench_similarity_f32.cpp

@@ -0,0 +1,118 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * 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.
+ */
+
+// Google Benchmark micro-benchmarks for the fp32 vector similarity kernels:
+//   cosine_f32, dot_product_f32, euclidean_f32
+//
+// Parameterised over the realistic embedding dimensions used in production:
+//   128, 256, 512, 1024, 1536, 3072
+//
+// Build (requires google-benchmark installed or available via pkg-config):
+//   meson setup build && ninja -C build bench_simd_kernels
+//
+// Run:
+//   ./build/bench_simd_kernels [--benchmark_filter=<pattern>]
+
+#include <benchmark/benchmark.h>
+#include <cmath>
+#include <vector>
+
+#include "jvector_simd.h"
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+// Deterministic, non-zero float vector: avoids degenerate cosine=NaN cases.
+static std::vector<float> make_vec(size_t n, float seed)
+{
+    std::vector<float> v(n);
+    for (size_t i = 0; i < n; ++i) {
+        v[i] = seed * (1.0f + static_cast<float>(i % 7) * 0.13f);
+        if (i % 3 == 0) v[i] = -v[i];
+        v[i] += 0.5f;
+    }
+    return v;
+}
+
+// Benchmark sizes matching production embedding dimensions.
+static const std::vector<int64_t> kBenchSizes = {128, 256, 512, 1024, 1536, 3072};
+
+// ---------------------------------------------------------------------------
+// dot_product_f32
+// ---------------------------------------------------------------------------
+
+static void BM_dot_product_f32(benchmark::State& state)
+{
+    const size_t n = static_cast<size_t>(state.range(0));
+    auto a = make_vec(n, 0.7f);
+    auto b = make_vec(n, 1.3f);
+
+    for (auto _ : state) {
+        float result = dot_product_f32(a.data(), 0, b.data(), 0, n);
+        benchmark::DoNotOptimize(result);
+    }
+
+    state.SetItemsProcessed(state.iterations() * static_cast<int64_t>(n));
+    state.SetBytesProcessed(state.iterations() * static_cast<int64_t>(n) * 2 * sizeof(float));
+}
+BENCHMARK(BM_dot_product_f32)->ArgsProduct({kBenchSizes});
+
+// ---------------------------------------------------------------------------
+// euclidean_f32
+// ---------------------------------------------------------------------------
+
+static void BM_euclidean_f32(benchmark::State& state)
+{
+    const size_t n = static_cast<size_t>(state.range(0));
+    auto a = make_vec(n, 0.7f);
+    auto b = make_vec(n, 1.3f);
+
+    for (auto _ : state) {
+        float result = euclidean_f32(a.data(), 0, b.data(), 0, n);
+        benchmark::DoNotOptimize(result);
+    }
+
+    state.SetItemsProcessed(state.iterations() * static_cast<int64_t>(n));
+    state.SetBytesProcessed(state.iterations() * static_cast<int64_t>(n) * 2 * sizeof(float));
+}
+BENCHMARK(BM_euclidean_f32)->ArgsProduct({kBenchSizes});
+
+// ---------------------------------------------------------------------------
+// cosine_f32
+// ---------------------------------------------------------------------------
+
+static void BM_cosine_f32(benchmark::State& state)
+{
+    const size_t n = static_cast<size_t>(state.range(0));
+    auto a = make_vec(n, 0.7f);
+    auto b = make_vec(n, 1.3f);
+
+    for (auto _ : state) {
+        float result = cosine_f32(a.data(), 0, b.data(), 0, n);
+        benchmark::DoNotOptimize(result);
+    }
+
+    state.SetItemsProcessed(state.iterations() * static_cast<int64_t>(n));
+    state.SetBytesProcessed(state.iterations() * static_cast<int64_t>(n) * 2 * sizeof(float));
+}
+BENCHMARK(BM_cosine_f32)->ArgsProduct({kBenchSizes});
+
+// ---------------------------------------------------------------------------
+// Entry point — benchmark::Initialize parses --benchmark_* flags.
+// ---------------------------------------------------------------------------
+
+BENCHMARK_MAIN();

jvector-native/src/main/native/meson.build

@@ -54,7 +54,7 @@ isa_libs = []
 foreach isa : isa_variants
   lib = static_library(
     'simdKernels_' + isa['name'],
-    sources            : 'jvector_simd_kernels.cpp',
+    sources            : 'src/jvector_simd_kernels.cpp',
     include_directories: hwy_inc,
     cpp_args           : isa['args'] + ['-DJV_ISA=' + isa['namespace'], '-fvisibility=hidden']
   )
@@ -65,7 +65,7 @@ endforeach
 # set (AVX3 + VNNI, VBMI, VBMI2, IFMA, BITALG, VPOPCNTDQ, GFNI, VAES, VPCLMULQDQ).
 avx3_dl_lib = static_library(
   'simdKernels_avx3_dl',
-  sources            : 'jvector_avx3_dl_kernels.cpp',
+  sources            : 'src/jvector_avx3_dl_kernels.cpp',
   include_directories: hwy_inc,
   cpp_args           : [
     '-march=icelake-server',
@@ -81,7 +81,7 @@ isa_libs += avx3_dl_lib
 # Requires GCC >= 12 or Clang >= 14.
 avx3_spr_lib = static_library(
   'simdKernels_avx3_spr',
-  sources            : 'jvector_avx3_spr_kernels.cpp',
+  sources            : 'src/jvector_avx3_spr_kernels.cpp',
   include_directories: hwy_inc,
   cpp_args           : [
     '-march=sapphirerapids',
@@ -100,9 +100,9 @@ isa_libs += avx3_spr_lib
 # projects that manage their own dispatch (google/highway#1935).
 vectorutil_lib = shared_library(
   'jvector',
-  sources            : ['jvector_simd.cpp',
+  sources            : ['src/jvector_simd.cpp',
                         'third_party/highway/hwy/abort.cc'],
-  include_directories: [include_directories('.'), hwy_inc],
+  include_directories: [include_directories('src'), hwy_inc],
   cpp_args           : ['-DJVECTOR_BUILD', '-fvisibility=hidden'],
   link_whole         : isa_libs,
   version            : meson.project_version(),
@@ -112,7 +112,7 @@ vectorutil_lib = shared_library(
 # Dependency object for use by executables/tests in this build tree.
 vectorutil_dep = declare_dependency(
   link_with          : vectorutil_lib,
-  include_directories: include_directories('.'),
+  include_directories: include_directories('src'),
 )
 
 ## Example driver that exercises the runtime-dispatch API.
@@ -121,40 +121,33 @@ vectorutil_dep = declare_dependency(
 #  sources     : 'examples/cpp_driver.cpp',
 #  dependencies: vectorutil_dep,
 #)
-#
-## ---- Tests -----------------------------------------------------------------
-#gtest_dep = dependency('gtest_main', required: true)
-#
-#test_exe = executable(
-#  'test_kernels',
-#  sources     : [
-#    'tests/test_kernels.cpp',
-#    'tests/test_cpuFeatures.cpp',
-#  ],
-#  dependencies: [vectorutil_dep, gtest_dep],
-#)
-#
-#test('kernels',     test_exe, protocol: 'gtest', suite: 'kernels')
-#test('cpu_features', test_exe, protocol: 'gtest', suite: 'cpu',
-#     args: ['--gtest_filter=CpuFeaturesTest.*'])
-#
-## ---- Benchmarks ------------------------------------------------------------
-#gbench_dep = dependency('benchmark', required: false)
-#if gbench_dep.found()
-#  executable(
-#    'bench_kernels',
-#    sources     : 'benchmarks/bench_kernels.cpp',
-#    dependencies: [vectorutil_dep, gbench_dep],
-#    cpp_args    : ['-O3'],
-#  )
-#endif
-#
-#rust_enabled = add_languages('rust', required: false)
-#if rust_enabled
-#  executable(
-#    'rust_driver',
-#    'examples/rust_driver.rs',
-#    link_with: vectorutil_lib,
-#  )
-#endif
-#
+
+# ---- Tests -----------------------------------------------------------------
+gtest_dep = dependency('gtest_main', required: false)
+
+if gtest_dep.found()
+  simd_kernels_test = executable(
+    'test_simd_kernels',
+    sources     : [
+      'tests/test_helpers.cpp',
+      'tests/test_similarity.cpp',
+      'tests/test_elementwise.cpp',
+      'tests/test_cpu_features.cpp',
+    ],
+    dependencies: [vectorutil_dep, gtest_dep],
+  )
+
+  test('simd_kernels', simd_kernels_test, protocol: 'gtest', suite: 'simd_kernels')
+
+endif
+# ---- Benchmarks ------------------------------------------------------------
+gbench_dep = dependency('benchmark', required: false)
+
+if gbench_dep.found()
+  executable(
+    'bench_simd_kernels',
+    sources     : 'benchmarks/bench_similarity_f32.cpp',
+    dependencies: [vectorutil_dep, gbench_dep],
+    cpp_args    : ['-O3'],
+  )
+endif

jvector-native/src/main/native/jextract_vector_simd.sh → jvector-native/src/main/native/src/jextract_vector_simd.sh

@@ -2,6 +2,8 @@
 
 # fail on error
 set -e
+# print commands as they are executed
+set +x
 
 # Copyright DataStax, Inc.
 #
@@ -17,6 +19,22 @@ set -e
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# ---------------------------------------------------------------------------
+# Path anchors — all derived from the git repository root so the script works
+# regardless of the working directory it is invoked from (Maven sets
+# workingDirectory to the src directory, but developers may run it from
+# anywhere inside the repo).
+# ---------------------------------------------------------------------------
+REPO_ROOT="$(git -C "$(dirname "$0")" rev-parse --show-toplevel)"
+SCRIPT_DIR="${REPO_ROOT}/jvector-native/src/main/native/src"
+NATIVE_DIR="${REPO_ROOT}/jvector-native/src/main/native"
+MODULE_ROOT="${REPO_ROOT}/jvector-native"
+
+HIGHWAY_DIR="${NATIVE_DIR}/third_party/highway"
+BUILD_DIR="${MODULE_ROOT}/target/meson-build"
+RESOURCES_DIR="${MODULE_ROOT}/src/main/resources"
+JAVA_OUT_DIR="${MODULE_ROOT}/src/main/java"
+
 if [ "$1" == "--auto-install-deps" ] ; then AUTO_INSTALL_DEPS=true ; shift ; fi
 printf "AUTO_INSTALL_DEPS=%s\n" "${AUTO_INSTALL_DEPS}"
 
@@ -29,13 +47,13 @@ if [ "$BUILDTYPE" != "release" ] && [ "$BUILDTYPE" != "debug" ] && [ "$BUILDTYPE
 fi
 printf "BUILDTYPE=%s\n" "${BUILDTYPE}"
 
-mkdir -p ../resources
+mkdir -p "${RESOURCES_DIR}"
+
 # compile jvector_simd_check.cpp as x86-64
 # compile jvector_simd.cpp as skylake-avx512
 # produce one shared library
 
 # Check that the Google Highway submodule has been initialised
-HIGHWAY_DIR="third_party/highway"
 if [ ! -f "${HIGHWAY_DIR}/hwy/highway.h" ]; then
   echo "ERROR: Google Highway submodule not found at ${HIGHWAY_DIR}."
   echo "       Run the following command from the repository root to fix this:"
@@ -80,24 +98,23 @@ if [ "$(printf '%s\n' "$MIN_GCC_VERSION" "$CURRENT_GPP_VERSION" | sort -V | head
     exit 1
 fi
 
-BUILD_DIR="../../../target/meson-build"
-rm -rf ../resources/libjvector.so
+rm -rf "${RESOURCES_DIR}/libjvector.so"
 
 # Configure (--wipe resets any stale configuration) then compile
-meson setup "${BUILD_DIR}" \
+meson setup "${BUILD_DIR}" "${NATIVE_DIR}" \
     --wipe \
     --buildtype="${BUILDTYPE}"
 
 meson compile -C "${BUILD_DIR}"
 
 # The versioned .so (e.g. libjvector.so.0.1.0) is the real file; symlinks point to it.
-# Copy it to ../resources/ as the plain libjvector.so for Java System.load().
+# Copy it to src/main/resources/ so Maven packages it into the jar for LibraryLoader.
 SOFILE=$(find "${BUILD_DIR}" -maxdepth 1 -name 'libjvector.so.*' -type f | head -1)
 if [ -z "${SOFILE}" ]; then
     echo "ERROR: libjvector.so not found in ${BUILD_DIR} after build."
     exit 1
 fi
-cp "${SOFILE}" ../resources/libjvector.so
+cp "${SOFILE}" "${RESOURCES_DIR}/libjvector.so"
 
 # Generate Java source code
 # Should only be run when c header changes
@@ -109,11 +126,12 @@ then
 fi
 
 jextract \
-  --output ../java \
+  --output "${JAVA_OUT_DIR}" \
   -t io.github.jbellis.jvector.vector.cnative \
-  -I . \
+  -I "${SCRIPT_DIR}" \
   --header-class-name NativeSimdOps \
-  jvector_simd.h
+  "${SCRIPT_DIR}/jvector_simd.h"
 
 # Set critical linker option with heap-based segments for all generated methods
-sed -i 's/DESC)/DESC, Linker.Option.critical(true))/g' ../java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java
+sed -i 's/DESC)/DESC, Linker.Option.critical(true))/g' \
+  "${JAVA_OUT_DIR}/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java"