Add functionality to test if the host CPU supports native SIMD instructions (#107429)

libs/native/libraries/build.gradle

@@ -18,7 +18,7 @@ configurations {
 }
 
 var zstdVersion = "1.5.5"
-var vecVersion = "1.0.1"
+var vecVersion = "1.0.3"
 
 repositories {
   exclusiveContent {

libs/native/src/main/java/org/elasticsearch/nativeaccess/PosixNativeAccess.java

@@ -27,7 +27,7 @@ abstract class PosixNativeAccess extends AbstractNativeAccess {
 
     static VectorSimilarityFunctions vectorSimilarityFunctionsOrNull(NativeLibraryProvider libraryProvider) {
         if (isNativeVectorLibSupported()) {
-            var lib = new VectorSimilarityFunctions(libraryProvider.getLibrary(VectorLibrary.class));
+            var lib = libraryProvider.getLibrary(VectorLibrary.class).getVectorSimilarityFunctions();
             logger.info("Using native vector library; to disable start with -D" + ENABLE_JDK_VECTOR_LIBRARY + "=false");
             return lib;
         }

libs/native/src/main/java/org/elasticsearch/nativeaccess/VectorSimilarityFunctions.java

@@ -8,25 +8,16 @@
 
 package org.elasticsearch.nativeaccess;
 
-import org.elasticsearch.nativeaccess.lib.VectorLibrary;
-
 import java.lang.invoke.MethodHandle;
 
 /**
- * Utility class providing vector similarity functions.
+ * Utility interface providing vector similarity functions.
  *
  * <p> MethodHandles are returned to avoid a static reference to MemorySegment,
  * which is not in the currently lowest compile version, JDK 17. Code consuming
  * the method handles will, by definition, require access to MemorySegment.
  */
-public final class VectorSimilarityFunctions implements VectorLibrary {
-
-    private final VectorLibrary vectorLibrary;
-
-    VectorSimilarityFunctions(VectorLibrary vectorLibrary) {
-        this.vectorLibrary = vectorLibrary;
-    }
-
+public interface VectorSimilarityFunctions {
     /**
      * Produces a method handle returning the dot product of byte (signed int8) vectors.
      *
@@ -34,9 +25,7 @@ public final class VectorSimilarityFunctions implements VectorLibrary {
      * its first and second arguments will be {@code MemorySegment}, whose contents is the
      * vector data bytes. The third argument is the length of the vector data.
      */
-    public MethodHandle dotProductHandle() {
-        return vectorLibrary.dotProductHandle();
-    }
+    MethodHandle dotProductHandle();
 
     /**
      * Produces a method handle returning the square distance of byte (signed int8) vectors.
@@ -45,7 +34,5 @@ public final class VectorSimilarityFunctions implements VectorLibrary {
      * its first and second arguments will be {@code MemorySegment}, whose contents is the
      * vector data bytes. The third argument is the length of the vector data.
      */
-    public MethodHandle squareDistanceHandle() {
-        return vectorLibrary.squareDistanceHandle();
-    }
+    MethodHandle squareDistanceHandle();
 }

libs/native/src/main/java/org/elasticsearch/nativeaccess/lib/VectorLibrary.java

@@ -8,7 +8,8 @@
 
 package org.elasticsearch.nativeaccess.lib;
 
-import java.lang.invoke.MethodHandle;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.nativeaccess.VectorSimilarityFunctions;
 
 /**
  * A VectorLibrary is just an adaptation of the factory for a NativeLibrary.
@@ -16,8 +17,6 @@ import java.lang.invoke.MethodHandle;
  * for native implementations.
  */
 public non-sealed interface VectorLibrary extends NativeLibrary {
-
-    MethodHandle dotProductHandle();
-
-    MethodHandle squareDistanceHandle();
+    @Nullable
+    VectorSimilarityFunctions getVectorSimilarityFunctions();
 }

libs/native/src/main21/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java

@@ -8,6 +8,7 @@
 
 package org.elasticsearch.nativeaccess.jdk;
 
+import org.elasticsearch.nativeaccess.VectorSimilarityFunctions;
 import org.elasticsearch.nativeaccess.lib.VectorLibrary;
 
 import java.lang.foreign.FunctionDescriptor;
@@ -23,12 +24,33 @@ import static org.elasticsearch.nativeaccess.jdk.LinkerHelper.downcallHandle;
 
 public final class JdkVectorLibrary implements VectorLibrary {
 
+    static final VectorSimilarityFunctions INSTANCE;
+
     static {
         System.loadLibrary("vec");
+        final MethodHandle vecCaps$mh = downcallHandle("vec_caps", FunctionDescriptor.of(JAVA_INT));
+
+        try {
+            int caps = (int) vecCaps$mh.invokeExact();
+            if (caps != 0) {
+                INSTANCE = new JdkVectorSimilarityFunctions();
+            } else {
+                INSTANCE = null;
+            }
+        } catch (Throwable t) {
+            throw new AssertionError(t);
+        }
     }
 
     public JdkVectorLibrary() {}
 
+    @Override
+    public VectorSimilarityFunctions getVectorSimilarityFunctions() {
+        return INSTANCE;
+    }
+
+    private static final class JdkVectorSimilarityFunctions implements VectorSimilarityFunctions {
+
         static final MethodHandle dot8stride$mh = downcallHandle("dot8s_stride", FunctionDescriptor.of(JAVA_INT));
         static final MethodHandle sqr8stride$mh = downcallHandle("sqr8s_stride", FunctionDescriptor.of(JAVA_INT));
 
@@ -49,6 +71,7 @@ public final class JdkVectorLibrary implements VectorLibrary {
 
         /**
          * Computes the dot product of given byte vectors.
+         *
          * @param a      address of the first vector
          * @param b      address of the second vector
          * @param length the vector dimensions
@@ -78,6 +101,7 @@ public final class JdkVectorLibrary implements VectorLibrary {
 
         /**
          * Computes the square distance of given byte vectors.
+         *
          * @param a      address of the first vector
          * @param b      address of the second vector
          * @param length the vector dimensions
@@ -145,8 +169,8 @@ public final class JdkVectorLibrary implements VectorLibrary {
             try {
                 var lookup = MethodHandles.lookup();
                 var mt = MethodType.methodType(int.class, MemorySegment.class, MemorySegment.class, int.class);
-            DOT_HANDLE = lookup.findStatic(JdkVectorLibrary.class, "dotProduct", mt);
-            SQR_HANDLE = lookup.findStatic(JdkVectorLibrary.class, "squareDistance", mt);
+                DOT_HANDLE = lookup.findStatic(JdkVectorSimilarityFunctions.class, "dotProduct", mt);
+                SQR_HANDLE = lookup.findStatic(JdkVectorSimilarityFunctions.class, "squareDistance", mt);
             } catch (NoSuchMethodException | IllegalAccessException e) {
                 throw new RuntimeException(e);
             }
@@ -161,4 +185,5 @@ public final class JdkVectorLibrary implements VectorLibrary {
         public MethodHandle squareDistanceHandle() {
             return SQR_HANDLE;
         }
+    }
 }

libs/vec/native/publish_vec_binaries.sh

@@ -19,7 +19,7 @@ if [ -z "$ARTIFACTORY_API_KEY" ]; then
   exit 1;
 fi
 
-VERSION="1.0.1"
+VERSION="1.0.3"
 ARTIFACTORY_REPOSITORY="${ARTIFACTORY_REPOSITORY:-https://artifactory.elastic.dev/artifactory/elasticsearch-native/}"
 TEMP=$(mktemp -d)
 

libs/vec/native/src/vec/c/vec.c

@@ -18,6 +18,34 @@
 #define SQR8S_STRIDE_BYTES_LEN 16
 #endif
 
+#ifdef __linux__
+    #include <sys/auxv.h>
+    #include <asm/hwcap.h>
+    #ifndef HWCAP_NEON
+    #define HWCAP_NEON 0x1000
+    #endif
+#endif
+
+#ifdef __APPLE__
+#include <TargetConditionals.h>
+#endif
+
+EXPORT int vec_caps() {
+#ifdef __APPLE__
+    #ifdef TARGET_OS_OSX
+        // All M series Apple silicon support Neon instructions
+        return 1;
+    #else
+        #error "Unsupported Apple platform"
+    #endif
+#elif __linux__
+    int hwcap = getauxval(AT_HWCAP);
+    return (hwcap & HWCAP_NEON) != 0;
+#else
+    #error "Unsupported aarch64 platform"
+#endif
+}
+
 EXPORT int dot8s_stride() {
     return DOT8_STRIDE_BYTES_LEN;
 }

libs/vec/native/src/vec/headers/vec.h

@@ -6,7 +6,15 @@
  * Side Public License, v 1.
  */
 
+#ifdef _MSC_VER
+#define EXPORT extern "C" __declspec(dllexport)
+#elif defined(__GNUC__) && !defined(__clang__)
 #define EXPORT __attribute__((externally_visible,visibility("default")))
+#elif __clang__
+#define EXPORT __attribute__((visibility("default")))
+#endif
+
+EXPORT int vec_caps();
 
 EXPORT int dot8s_stride();