Allow incubating Panama Vector in simdvec, and add vectorized ipByteBin (#112933)

Add support for vectorized ipByteBin.

The structure of the implementation and loading framework mirror that of Lucene, but is simplified by avoiding reflective loading since ES has support for a MRJar section for 21.

For now, we just disable warnings-as-errors in this small sourceset, since -Xlint:-incubating is only support since JDK 22. The number of source files is small here. Will investigate how to assert that just the single incubating warning is emitted by javac, at a later point.

Author: ChrisHegarty

docs/changelog/112933.yaml

@@ -0,0 +1,5 @@
+pr: 112933
+summary: "Allow incubating Panama Vector in simdvec, and add vectorized `ipByteBin`"
+area: Search
+type: enhancement
+issues: []

libs/simdvec/build.gradle

@@ -23,6 +23,19 @@ dependencies {
   }
 }
 
+tasks.named("compileMain21Java").configure {
+  options.compilerArgs << '--add-modules=jdk.incubator.vector'
+  // we remove Werror, since incubating suppression (-Xlint:-incubating)
+  // is only support since JDK 22
+  options.compilerArgs -= '-Werror'
+}
+
+test {
+  if (JavaVersion.current().majorVersion.toInteger() >= 21) {
+    jvmArgs '--add-modules=jdk.incubator.vector'
+  }
+}
+
 tasks.withType(CheckForbiddenApisTask).configureEach {
   replaceSignatureFiles 'jdk-signatures'
 }

libs/simdvec/src/main/java/module-info.java

@@ -10,6 +10,7 @@
 module org.elasticsearch.simdvec {
     requires org.elasticsearch.nativeaccess;
     requires org.apache.lucene.core;
+    requires org.elasticsearch.logging;
 
     exports org.elasticsearch.simdvec to org.elasticsearch.server;
 }

libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.simdvec;
+
+import org.elasticsearch.simdvec.internal.vectorization.ESVectorUtilSupport;
+import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider;
+
+import static org.elasticsearch.simdvec.internal.vectorization.ESVectorUtilSupport.B_QUERY;
+
+public class ESVectorUtil {
+
+    private static final ESVectorUtilSupport IMPL = ESVectorizationProvider.getInstance().getVectorUtilSupport();
+
+    public static long ipByteBinByte(byte[] q, byte[] d) {
+        if (q.length != d.length * B_QUERY) {
+            throw new IllegalArgumentException("vector dimensions incompatible: " + q.length + "!= " + B_QUERY + " x " + d.length);
+        }
+        return IMPL.ipByteBinByte(q, d);
+    }
+}

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.simdvec.internal.vectorization;
+
+import org.apache.lucene.util.BitUtil;
+
+final class DefaultESVectorUtilSupport implements ESVectorUtilSupport {
+
+    DefaultESVectorUtilSupport() {}
+
+    @Override
+    public long ipByteBinByte(byte[] q, byte[] d) {
+        return ipByteBinByteImpl(q, d);
+    }
+
+    public static long ipByteBinByteImpl(byte[] q, byte[] d) {
+        long ret = 0;
+        int size = d.length;
+        for (int i = 0; i < B_QUERY; i++) {
+            int r = 0;
+            long subRet = 0;
+            for (final int upperBound = d.length & -Integer.BYTES; r < upperBound; r += Integer.BYTES) {
+                subRet += Integer.bitCount((int) BitUtil.VH_NATIVE_INT.get(q, i * size + r) & (int) BitUtil.VH_NATIVE_INT.get(d, r));
+            }
+            for (; r < d.length; r++) {
+                subRet += Integer.bitCount((q[i * size + r] & d[r]) & 0xFF);
+            }
+            ret += subRet << i;
+        }
+        return ret;
+    }
+}

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorizationProvider.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.simdvec.internal.vectorization;
+
+final class DefaultESVectorizationProvider extends ESVectorizationProvider {
+    private final ESVectorUtilSupport vectorUtilSupport;
+
+    DefaultESVectorizationProvider() {
+        vectorUtilSupport = new DefaultESVectorUtilSupport();
+    }
+
+    @Override
+    public ESVectorUtilSupport getVectorUtilSupport() {
+        return vectorUtilSupport;
+    }
+}

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java

@@ -0,0 +1,17 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.simdvec.internal.vectorization;
+
+public interface ESVectorUtilSupport {
+
+    short B_QUERY = 4;
+
+    long ipByteBinByte(byte[] q, byte[] d);
+}

libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.simdvec.internal.vectorization;
+
+import java.util.Objects;
+
+public abstract class ESVectorizationProvider {
+
+    public static ESVectorizationProvider getInstance() {
+        return Objects.requireNonNull(
+            ESVectorizationProvider.Holder.INSTANCE,
+            "call to getInstance() from subclass of VectorizationProvider"
+        );
+    }
+
+    ESVectorizationProvider() {}
+
+    public abstract ESVectorUtilSupport getVectorUtilSupport();
+
+    // visible for tests
+    static ESVectorizationProvider lookup(boolean testMode) {
+        return new DefaultESVectorizationProvider();
+    }
+
+    /** This static holder class prevents classloading deadlock. */
+    private static final class Holder {
+        private Holder() {}
+
+        static final ESVectorizationProvider INSTANCE = lookup(false);
+    }
+}

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorizationProvider.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.simdvec.internal.vectorization;
+
+import org.apache.lucene.util.Constants;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
+
+import java.util.Locale;
+import java.util.Objects;
+import java.util.Optional;
+
+public abstract class ESVectorizationProvider {
+
+    protected static final Logger logger = LogManager.getLogger(ESVectorizationProvider.class);
+
+    public static ESVectorizationProvider getInstance() {
+        return Objects.requireNonNull(
+            ESVectorizationProvider.Holder.INSTANCE,
+            "call to getInstance() from subclass of VectorizationProvider"
+        );
+    }
+
+    ESVectorizationProvider() {}
+
+    public abstract ESVectorUtilSupport getVectorUtilSupport();
+
+    // visible for tests
+    static ESVectorizationProvider lookup(boolean testMode) {
+        final int runtimeVersion = Runtime.version().feature();
+        assert runtimeVersion >= 21;
+        if (runtimeVersion <= 23) {
+            // only use vector module with Hotspot VM
+            if (Constants.IS_HOTSPOT_VM == false) {
+                logger.warn("Java runtime is not using Hotspot VM; Java vector incubator API can't be enabled.");
+                return new DefaultESVectorizationProvider();
+            }
+            // is the incubator module present and readable (JVM providers may to exclude them or it is
+            // build with jlink)
+            final var vectorMod = lookupVectorModule();
+            if (vectorMod.isEmpty()) {
+                logger.warn(
+                    "Java vector incubator module is not readable. "
+                        + "For optimal vector performance, pass '--add-modules jdk.incubator.vector' to enable Vector API."
+                );
+                return new DefaultESVectorizationProvider();
+            }
+            vectorMod.ifPresent(ESVectorizationProvider.class.getModule()::addReads);
+            var impl = new PanamaESVectorizationProvider();
+            logger.info(
+                String.format(
+                    Locale.ENGLISH,
+                    "Java vector incubator API enabled; uses preferredBitSize=%d",
+                    PanamaESVectorUtilSupport.VECTOR_BITSIZE
+                )
+            );
+            return impl;
+        } else {
+            logger.warn(
+                "You are running with unsupported Java "
+                    + runtimeVersion
+                    + ". To make full use of the Vector API, please update Elasticsearch."
+            );
+        }
+        return new DefaultESVectorizationProvider();
+    }
+
+    private static Optional<Module> lookupVectorModule() {
+        return Optional.ofNullable(ESVectorizationProvider.class.getModule().getLayer())
+            .orElse(ModuleLayer.boot())
+            .findModule("jdk.incubator.vector");
+    }
+
+    /** This static holder class prevents classloading deadlock. */
+    private static final class Holder {
+        private Holder() {}
+
+        static final ESVectorizationProvider INSTANCE = lookup(false);
+    }
+}