Integrate the Incubating Panama Vector API

gradle/java/memorysegment-mrjar.gradle

@@ -40,6 +40,7 @@ configure(project(":lucene:core")) {
           "-Xlint:-options",
           "--patch-module", "java.base=${apijar}",
           "--add-exports", "java.base/java.lang.foreign=ALL-UNNAMED",
+          "--add-exports", "java.base/jdk.incubator.vector=ALL-UNNAMED", // This is a hack, but does it matter
         ]
       }
     }

gradle/testing/defaults-tests.gradle

@@ -122,7 +122,7 @@ allprojects {
 
       // Lucene needs to optional modules at runtime, which we want to enforce for testing
       // (if the runner JVM does not support them, it will fail tests):
-      jvmArgs '--add-modules', 'jdk.unsupported,jdk.management'
+      jvmArgs '--add-modules', 'jdk.unsupported,jdk.management,jdk.incubator.vector'
 
       def loggingConfigFile = layout.projectDirectory.file("${resources}/logging.properties")
       def tempDir = layout.projectDirectory.dir(testsTmpDir.toString())

lucene/core/src/java/org/apache/lucene/util/VectorUtil.java

@@ -17,9 +17,19 @@
 
 package org.apache.lucene.util;
 
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.security.AccessController;
+import java.security.PrivilegedAction;
+import java.util.logging.Logger;
+
 /** Utilities for computations with numeric arrays */
 public final class VectorUtil {
 
+  private static final Logger LOG = Logger.getLogger(VectorUtil.class.getName());
+
+  private static final VectorUtilProvider PROVIDER;
+
   private VectorUtil() {}
 
   /**
@@ -28,71 +38,7 @@ private VectorUtil() {}
    * @throws IllegalArgumentException if the vectors' dimensions differ.
    */
   public static float dotProduct(float[] a, float[] b) {
-    if (a.length != b.length) {
-      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
-    }
-    float res = 0f;
-    /*
-     * If length of vector is larger than 8, we use unrolled dot product to accelerate the
-     * calculation.
-     */
-    int i;
-    for (i = 0; i < a.length % 8; i++) {
-      res += b[i] * a[i];
-    }
-    if (a.length < 8) {
-      return res;
-    }
-    for (; i + 31 < a.length; i += 32) {
-      res +=
-          b[i + 0] * a[i + 0]
-              + b[i + 1] * a[i + 1]
-              + b[i + 2] * a[i + 2]
-              + b[i + 3] * a[i + 3]
-              + b[i + 4] * a[i + 4]
-              + b[i + 5] * a[i + 5]
-              + b[i + 6] * a[i + 6]
-              + b[i + 7] * a[i + 7];
-      res +=
-          b[i + 8] * a[i + 8]
-              + b[i + 9] * a[i + 9]
-              + b[i + 10] * a[i + 10]
-              + b[i + 11] * a[i + 11]
-              + b[i + 12] * a[i + 12]
-              + b[i + 13] * a[i + 13]
-              + b[i + 14] * a[i + 14]
-              + b[i + 15] * a[i + 15];
-      res +=
-          b[i + 16] * a[i + 16]
-              + b[i + 17] * a[i + 17]
-              + b[i + 18] * a[i + 18]
-              + b[i + 19] * a[i + 19]
-              + b[i + 20] * a[i + 20]
-              + b[i + 21] * a[i + 21]
-              + b[i + 22] * a[i + 22]
-              + b[i + 23] * a[i + 23];
-      res +=
-          b[i + 24] * a[i + 24]
-              + b[i + 25] * a[i + 25]
-              + b[i + 26] * a[i + 26]
-              + b[i + 27] * a[i + 27]
-              + b[i + 28] * a[i + 28]
-              + b[i + 29] * a[i + 29]
-              + b[i + 30] * a[i + 30]
-              + b[i + 31] * a[i + 31];
-    }
-    for (; i + 7 < a.length; i += 8) {
-      res +=
-          b[i + 0] * a[i + 0]
-              + b[i + 1] * a[i + 1]
-              + b[i + 2] * a[i + 2]
-              + b[i + 3] * a[i + 3]
-              + b[i + 4] * a[i + 4]
-              + b[i + 5] * a[i + 5]
-              + b[i + 6] * a[i + 6]
-              + b[i + 7] * a[i + 7];
-    }
-    return res;
+    return PROVIDER.dotProduct(a, b);
   }
 
   /**
@@ -270,4 +216,134 @@ public static float dotProductScore(byte[] a, byte[] b) {
     float denom = (float) (a.length * (1 << 15));
     return 0.5f + dotProduct(a, b) / denom;
   }
+
+  interface VectorUtilProvider {
+
+    // just dot product for now
+    float dotProduct(float[] a, float[] b);
+  }
+
+  private static VectorUtilProvider lookupProvider() {
+    // TODO: add a check
+    final int runtimeVersion = Runtime.version().feature();
+    if (runtimeVersion == 20) { // TODO: do we want JDK 19?
+      try {
+        ensureReadability();
+        final var lookup = MethodHandles.lookup();
+        final var cls = lookup.findClass("org.apache.lucene.util.JDKVectorUtilProvider");
+        // we use method handles, so we do not need to deal with setAccessible as we have private
+        // access through the lookup:
+        final var constr = lookup.findConstructor(cls, MethodType.methodType(void.class));
+        try {
+          return (VectorUtilProvider) constr.invoke();
+        } catch (RuntimeException | Error e) {
+          throw e;
+        } catch (Throwable th) {
+          throw new AssertionError(th);
+        }
+      } catch (NoSuchMethodException | IllegalAccessException e) {
+        throw new LinkageError("JDKVectorUtilProvider is missing correctly typed constructor", e);
+      } catch (ClassNotFoundException cnfe) {
+        throw new LinkageError("JDKVectorUtilProvider is missing in Lucene JAR file", cnfe);
+      }
+    } else if (runtimeVersion >= 21) {
+      LOG.warning(
+          "You are running with Java 21 or later. To make full use of the Vector API, please update Apache Lucene.");
+    }
+    return new LuceneVectorUtilProvider();
+  }
+
+  // Extracted to a method to be able to apply the SuppressForbidden annotation
+  @SuppressWarnings("removal")
+  @SuppressForbidden(reason = "security manager")
+  private static <T> T doPrivileged(PrivilegedAction<T> action) {
+    return AccessController.doPrivileged(action);
+  }
+
+  static void ensureReadability() {
+    ModuleLayer.boot().modules().stream()
+        .filter(m -> m.getName().equals("jdk.incubator.vector"))
+        .findFirst()
+        .ifPresentOrElse(
+            vecMod -> VectorUtilProvider.class.getModule().addReads(vecMod),
+            () -> LOG.warning("vector incubator module not present"));
+  }
+
+  static {
+    PROVIDER =
+        doPrivileged(VectorUtil::lookupProvider); // TODO check what permissions we actually need
+  }
+
+  static final class LuceneVectorUtilProvider implements VectorUtilProvider {
+
+    @Override
+    public float dotProduct(float[] a, float[] b) {
+      if (a.length != b.length) {
+        throw new IllegalArgumentException(
+            "vector dimensions differ: " + a.length + "!=" + b.length);
+      }
+      float res = 0f;
+      /*
+       * If length of vector is larger than 8, we use unrolled dot product to accelerate the
+       * calculation.
+       */
+      int i;
+      for (i = 0; i < a.length % 8; i++) {
+        res += b[i] * a[i];
+      }
+      if (a.length < 8) {
+        return res;
+      }
+      for (; i + 31 < a.length; i += 32) {
+        res +=
+            b[i + 0] * a[i + 0]
+                + b[i + 1] * a[i + 1]
+                + b[i + 2] * a[i + 2]
+                + b[i + 3] * a[i + 3]
+                + b[i + 4] * a[i + 4]
+                + b[i + 5] * a[i + 5]
+                + b[i + 6] * a[i + 6]
+                + b[i + 7] * a[i + 7];
+        res +=
+            b[i + 8] * a[i + 8]
+                + b[i + 9] * a[i + 9]
+                + b[i + 10] * a[i + 10]
+                + b[i + 11] * a[i + 11]
+                + b[i + 12] * a[i + 12]
+                + b[i + 13] * a[i + 13]
+                + b[i + 14] * a[i + 14]
+                + b[i + 15] * a[i + 15];
+        res +=
+            b[i + 16] * a[i + 16]
+                + b[i + 17] * a[i + 17]
+                + b[i + 18] * a[i + 18]
+                + b[i + 19] * a[i + 19]
+                + b[i + 20] * a[i + 20]
+                + b[i + 21] * a[i + 21]
+                + b[i + 22] * a[i + 22]
+                + b[i + 23] * a[i + 23];
+        res +=
+            b[i + 24] * a[i + 24]
+                + b[i + 25] * a[i + 25]
+                + b[i + 26] * a[i + 26]
+                + b[i + 27] * a[i + 27]
+                + b[i + 28] * a[i + 28]
+                + b[i + 29] * a[i + 29]
+                + b[i + 30] * a[i + 30]
+                + b[i + 31] * a[i + 31];
+      }
+      for (; i + 7 < a.length; i += 8) {
+        res +=
+            b[i + 0] * a[i + 0]
+                + b[i + 1] * a[i + 1]
+                + b[i + 2] * a[i + 2]
+                + b[i + 3] * a[i + 3]
+                + b[i + 4] * a[i + 4]
+                + b[i + 5] * a[i + 5]
+                + b[i + 6] * a[i + 6]
+                + b[i + 7] * a[i + 7];
+      }
+      return res;
+    }
+  }
 }

lucene/core/src/java20/org/apache/lucene/util/JDKVectorUtilProvider.java

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+package org.apache.lucene.util;
+
+import jdk.incubator.vector.FloatVector;
+import jdk.incubator.vector.VectorOperators;
+import jdk.incubator.vector.VectorSpecies;
+
+public final class JDKVectorUtilProvider implements VectorUtil.VectorUtilProvider {
+
+  static final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_PREFERRED;
+
+  public JDKVectorUtilProvider() {}
+
+  @Override
+  public float dotProduct(float[] a, float[] b) {
+    var upperBound = SPECIES.loopBound(a.length);
+    var sum = FloatVector.zero(SPECIES);
+    var i = 0;
+    for (; i < upperBound; i += SPECIES.length()) {
+      var va = FloatVector.fromArray(SPECIES, a, i);
+      var vb = FloatVector.fromArray(SPECIES, b, i);
+      sum = va.fma(vb, sum);
+    }
+    var c = sum.reduceLanes(VectorOperators.ADD);
+    for (; i < a.length; i++) { // tail loop
+      c += a[i] * b[i];
+    }
+    return c;
+  }
+}