[GPU] Optimize merge memory usage

distribution/tools/server-cli/src/main/java/org/elasticsearch/server/cli/SystemJvmOptions.java

@@ -60,6 +60,9 @@ static List<String> systemJvmOptions(Settings nodeSettings, final Map<String, St
                 "-Dio.netty.noUnsafe=true",
                 "-Dio.netty.noKeySetOptimization=true",
                 "-Dio.netty.recycler.maxCapacityPerThread=0",
+                // temporary until we get access to raw vectors in a future Lucene version
+                "--add-opens=org.apache.lucene.core/org.apache.lucene.codecs.lucene99=org.elasticsearch.server",
+                "--add-opens=org.apache.lucene.core/org.apache.lucene.internal.vectorization=org.elasticsearch.server",
                 // log4j 2
                 "-Dlog4j.shutdownHookEnabled=false",
                 "-Dlog4j2.disable.jmx=true",

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

@@ -0,0 +1,16 @@
+/*
+ * 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.apache.lucene.util.quantization.QuantizedByteVectorValues;
+
+public interface QuantizedByteVectorValuesAccess {
+    QuantizedByteVectorValues get();
+}

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/Int7SQVectorScorerSupplier.java

@@ -14,6 +14,7 @@
 import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
 import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
 import org.apache.lucene.util.quantization.ScalarQuantizedVectorSimilarity;
+import org.elasticsearch.simdvec.QuantizedByteVectorValuesAccess;
 
 import java.io.IOException;
 import java.lang.foreign.MemorySegment;
@@ -23,7 +24,7 @@
 import static org.apache.lucene.index.VectorSimilarityFunction.MAXIMUM_INNER_PRODUCT;
 import static org.apache.lucene.util.quantization.ScalarQuantizedVectorSimilarity.fromVectorSimilarity;
 
-public abstract sealed class Int7SQVectorScorerSupplier implements RandomVectorScorerSupplier {
+public abstract sealed class Int7SQVectorScorerSupplier implements RandomVectorScorerSupplier, QuantizedByteVectorValuesAccess {
 
     static final byte BITS = 7;
 
@@ -107,6 +108,11 @@ public void setScoringOrdinal(int node) throws IOException {
         };
     }
 
+    @Override
+    public QuantizedByteVectorValues get() {
+        return values;
+    }
+
     public static final class EuclideanSupplier extends Int7SQVectorScorerSupplier {
 
         public EuclideanSupplier(MemorySegmentAccessInput input, QuantizedByteVectorValues values, float scoreCorrectionConstant) {

server/src/main/java/module-info.java

@@ -490,6 +490,7 @@
     exports org.elasticsearch.lucene.util.automaton;
     exports org.elasticsearch.index.codec.perfield;
     exports org.elasticsearch.index.codec.vectors to org.elasticsearch.test.knn, org.elasticsearch.gpu;
+    exports org.elasticsearch.index.codec.vectors.reflect to org.elasticsearch.gpu;
     exports org.elasticsearch.index.codec.vectors.es818 to org.elasticsearch.test.knn;
     exports org.elasticsearch.inference.telemetry;
     exports org.elasticsearch.index.codec.vectors.diskbbq to org.elasticsearch.test.knn;

server/src/main/java/org/elasticsearch/index/codec/vectors/reflect/VectorsFormatReflectionUtils.java

@@ -0,0 +1,96 @@
+/*
+ * 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.index.codec.vectors.reflect;
+
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.util.hnsw.CloseableRandomVectorScorerSupplier;
+import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
+import org.elasticsearch.simdvec.QuantizedByteVectorValuesAccess;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+
+public class VectorsFormatReflectionUtils {
+    private static final VarHandle FLOAT_SUPPLIER_HANDLE;
+    private static final VarHandle BYTE_SUPPLIER_HANDLE;
+    private static final VarHandle FLOAT_VECTORS_HANDLE;
+
+    private static final Class<?> FLAT_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS;
+    private static final Class<?> SCALAR_QUANTIZED_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS;
+    private static final Class<?> FLOAT_SCORING_SUPPLIER_CLASS;
+    static {
+        try {
+            FLAT_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS = Class.forName(
+                "org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter$FlatCloseableRandomVectorScorerSupplier"
+            );
+            var lookup = MethodHandles.privateLookupIn(FLAT_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS, MethodHandles.lookup());
+            FLOAT_SUPPLIER_HANDLE = lookup.findVarHandle(
+                FLAT_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS,
+                "supplier",
+                RandomVectorScorerSupplier.class
+            );
+
+            FLOAT_SCORING_SUPPLIER_CLASS = Class.forName(
+                "org.apache.lucene.internal.vectorization.Lucene99MemorySegmentFloatVectorScorerSupplier"
+            );
+            lookup = MethodHandles.privateLookupIn(FLOAT_SCORING_SUPPLIER_CLASS, MethodHandles.lookup());
+            FLOAT_VECTORS_HANDLE = lookup.findVarHandle(FLOAT_SCORING_SUPPLIER_CLASS, "values", FloatVectorValues.class);
+
+            SCALAR_QUANTIZED_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS = Class.forName(
+                "org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter$ScalarQuantizedCloseableRandomVectorScorerSupplier"
+            );
+            lookup = MethodHandles.privateLookupIn(SCALAR_QUANTIZED_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS, MethodHandles.lookup());
+            BYTE_SUPPLIER_HANDLE = lookup.findVarHandle(
+                SCALAR_QUANTIZED_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS,
+                "supplier",
+                RandomVectorScorerSupplier.class
+            );
+
+        } catch (IllegalAccessException e) {
+            throw new AssertionError("should not happen, check opens", e);
+        } catch (ReflectiveOperationException e) {
+            throw new AssertionError(e);
+        }
+    }
+
+    public static RandomVectorScorerSupplier getFlatRandomVectorScorerInnerSupplier(CloseableRandomVectorScorerSupplier scorerSupplier) {
+        if (scorerSupplier.getClass().equals(FLAT_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS)) {
+            return (RandomVectorScorerSupplier) FLOAT_SUPPLIER_HANDLE.get(scorerSupplier);
+        }
+        return null;
+    }
+
+    public static RandomVectorScorerSupplier getScalarQuantizedRandomVectorScorerInnerSupplier(
+        CloseableRandomVectorScorerSupplier scorerSupplier
+    ) {
+        if (scorerSupplier.getClass().equals(SCALAR_QUANTIZED_CLOSEABLE_RANDOM_VECTOR_SCORER_SUPPLIER_CLASS)) {
+            return (RandomVectorScorerSupplier) BYTE_SUPPLIER_HANDLE.get(scorerSupplier);
+        }
+        return null;
+    }
+
+    public static HasIndexSlice getFloatScoringSupplierVectorOrNull(RandomVectorScorerSupplier scorerSupplier) {
+        if (FLOAT_SCORING_SUPPLIER_CLASS.isAssignableFrom(scorerSupplier.getClass())) {
+            var vectorValues = FLOAT_VECTORS_HANDLE.get(scorerSupplier);
+            if (vectorValues instanceof HasIndexSlice indexSlice) {
+                return indexSlice;
+            }
+        }
+        return null;
+    }
+
+    public static HasIndexSlice getByteScoringSupplierVectorOrNull(RandomVectorScorerSupplier scorerSupplier) {
+        if (scorerSupplier instanceof QuantizedByteVectorValuesAccess quantizedByteVectorValuesAccess) {
+            return quantizedByteVectorValuesAccess.get();
+        }
+        return null;
+    }
+}

x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUSupport.java

@@ -59,7 +59,7 @@ public static boolean isSupported(boolean logError) {
                     }
                 } else {
                     if (logError) {
-                        LOG.info("Found compatible GPU [{}] (id: [{}])", gpu.name(), gpu.gpuId());
+                        LOG.debug("Found compatible GPU [{}] (id: [{}])", gpu.name(), gpu.gpuId());
                     }
                     return true;
                 }