You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@lucene.apache.org by "shubhamvishu (via GitHub)" <gi...@apache.org> on 2025/12/17 03:37:59 UTC
[PR] Use native dot product in Lucene [lucene]
shubhamvishu opened a new pull request, #15508:
URL: https://github.com/apache/lucene/pull/15508
### Description
This is continuation of the work @goankur started in https://github.com/apache/lucene/pull/13572. Quoting the [comment](https://github.com/apache/lucene/pull/13572#issuecomment-3648627594) below from the initial PR.
> Here at Amazon (customer facing product search), we’ve been testing this native dot product implementation in our production environment(ARM - Graviton 2 and 3) and we see **5-14x** faster dot product computations in JMH benchmarks and we observed semantic latency improving from **62 msec** to **28 msec** (avg) for 4K embeddings(4.5 MM). Overall we saw **10-60%** improvement on end-end avg search latencies in different scenarios (different sized vectors, vector-focused search vs search combined with other workloads). We haven’t tested all other CPUs types yet. I'm working on a draft PR on top of this PR with following changes and planning to raise it soon :
>
> * Removing the overhead from heap to off-heap copying by utilizing `Linker.Option.critical`, which eliminates unnecessary copying
> * Runtime dispatch using `IFUNC` to choose SVE vs NEON vs scalar implementation at runtime based on available intrinsics
> * Build related changes to generate the binary
>
> We kept the native code isolated in the misc package and not getting it in the core module which we know is highly discouraged. Additionally, PR #15285 would later help eliminate some code duplication and enable a cleaner implementation similar to `PanamaVectorUtilSupport` - potentially through a `NativeVectorUtilSupport` class?
>
> Our benchmarking suggests substantial optimization potential for ARM-based deployments, and we believe this could benefit the broader Lucene community. We hope to make it easy for any Lucene user to opt-in to this alternative vector implementation ideally. We're committed to refining this implementation based on community feedback and addressing any concerns during the review process. I'm eager to hear the community's thoughts on this change, as there appears to be significant optimization potential for ARM architectures that could benefit many users. Thank you!
<!--
If this is your first contribution to Lucene, please make sure you have reviewed the contribution guide.
https://github.com/apache/lucene/blob/main/CONTRIBUTING.md
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2807500444
##########
lucene/CHANGES.txt:
##########
@@ -359,7 +357,7 @@ Optimizations
* GITHUB#15498: ExitableDirectoryReader keeps singleton SortedSetDocValues or SortedNumericDocValues as singletons. (Houston Putman)
-* GITHUB#15511, GITHUB#15696: Dynamic pruning for SORTED(_SET) fields with doc values skipper (Pan Guixin, Alan Woodward)
+* GITHUB#15511: Dynamic pruning for SORTED(_SET) fields with doc values skipper (Pan Guixin)
Review Comment:
Yeah, I fixed it. Now only shows this change in diff.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3879072807
Thank you @dweiss, really appreciate you taking the time to test it out!
So should we revert to the earlier approach (pre `NativeVectorUtilSupport` patch) we tried and stick to system property to override the behavior? Happy to make those changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2800876133
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
I agree, I'll soon push a version that is least intrusive.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3702126580
Please give me some time - I've been away on a winter break.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3878247802
Hi @shubhamvishu . I've tried to untangle this on my branch but I failed, sorry. My motivation was to make those alternative vector support implementations fairly independent but they're all so tightly coupled that I couldn't make any progress. Native delegates to panama, panama delegates to default... I'm not happy with the outcome of my refactoring and I suggest we roll back the changes you made (following my suggestions) to minimize the surface of this patch over current main.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3880995259
@dweiss I made those changes in the latest revision. Thank you!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3664141898
For the record, any Java library that ships with native libs has similar integration problems. You can look at jansi or any other lib that has native components for inspiration (and to get an idea) how hairy it becomes. Here's jansi's docker-based cross-compilation build.
https://github.com/fusesource/jansi/blob/master/Makefile
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2627993323
##########
lucene/core/build.gradle:
##########
@@ -31,6 +31,14 @@ dependencies {
moduleTestImplementation project(':lucene:test-framework')
}
+test {
+ dependsOn ':lucene:misc:buildNative'
+ systemProperty(
+ "java.library.path",
+ project(":lucene:misc").layout.buildDirectory.get().asFile.absolutePath + "/libs/dotProduct/shared"
+ )
+}
+
Review Comment:
Yeah, makes sense
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "uschindler (via GitHub)" <gi...@apache.org>.
uschindler commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2983184169
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
Review Comment:
this class should be package private
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorizationProvider.java:
##########
@@ -0,0 +1,79 @@
+/*
+ * 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.internal.vectorization;
+
+import java.io.IOException;
+import java.util.Locale;
+import java.util.logging.Logger;
+import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Constants;
+import org.apache.lucene.util.SuppressForbidden;
+
+/** Native provider returning native implementations which delegates to Panama implementation. */
+public final class NativeVectorizationProvider extends VectorizationProvider {
Review Comment:
this class should be package private
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2962527526
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
> I'd like to get rid of this magic. it slows down,.....
> :
> Ifwe don't want to do it statically ........
I like the suggestion, makes sense.
> I wonder why tests pass, do we really test all methods here?
We are not running it as part of the tests since we don't generate any binary (its still in the user court; user needs to pass the binary for tests to use it overriding whichever implementations they have) and the current [c code](https://github.com/apache/lucene/blob/main/lucene/misc/src/c/dotProduct.c) only have dotProduct impl and I tested few others but not all the existing VectorUtil methods. It would be good to always run the tests to ensure it catches all issues but idk if we want to go with generating a binary path. I'm thinking if its fine to not provide any platform specific binaries(abiding by what we discussed earlier) but just extend the c code to have basic VectorUtil implementations and test that on the unit test(which this PR was doing initially). Maybe you have some better ideas to how we can have better test coverage here with/without having binary? or just use fake(but simple like DefaultVectorUtil) java implementations to test each operation?
Right now this piece is bit disconnected due to no automatic test coverage running.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2962641734
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
I opened https://github.com/apache/lucene/issues/15840 to discuss the ideas on improving the test coverage for the native path in our CI workflows.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "rmuir (via GitHub)" <gi...@apache.org>.
rmuir commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3695034856
> Very Interesting! I'd love to explore this path and would be good in general if we could have some reliable way to provide efficient binaries to users. Maybe we could pursue this as a followup and keep this change just focussed on adding support to use the provided binary (next we could try adding actual binary support)?
Yes, there's no need to create binaries here. Was just a suggestion to avoid some hellacious network of Dockerfiles: for this kind of purpose, zig can be used as a better c compiler that "just works" to cross-compile all those binaries.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3667195378
I get the performance improvement is nice but it'll be difficult to maintain and test properly the way it's currently implemented. Not to mention most folks out there use Lucene without the possibility to recompile from sources....
I don't mind adding a native implementation but maybe this should be integrated via more standard mechanism (like service lookup) and then everything would live in that optional "native" module? This would open up the possibility to compile such a module independently. Sorry for my lack of enthusiasm about this but I can tell it'll be a problem to maintain even by looking at it (properties, etc.).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2807368129
##########
lucene/CHANGES.txt:
##########
@@ -359,7 +357,7 @@ Optimizations
* GITHUB#15498: ExitableDirectoryReader keeps singleton SortedSetDocValues or SortedNumericDocValues as singletons. (Houston Putman)
-* GITHUB#15511, GITHUB#15696: Dynamic pruning for SORTED(_SET) fields with doc values skipper (Pan Guixin, Alan Woodward)
+* GITHUB#15511: Dynamic pruning for SORTED(_SET) fields with doc values skipper (Pan Guixin)
Review Comment:
there are some differences here which look odd - perhaps this patch needs to be merged against main before it's put back in?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-4120417896
Thanks @uschindler ! I opened https://github.com/apache/lucene/pull/15866 to make the classes pkg-private and also fix the method @ExE-Boss pointed out above(we can eventually get rid of it with https://github.com/apache/lucene/issues/15840 as a followup). Looking for your thoughts and review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3673651373
@dweiss I tried to make it more pluggable and added the `NativeVectorUtilSupport` which let users opt-in to use native binaries on the fly without doing anything extra. If there is library present and also have the required implementation it automatically switches to native implementation or else fallback to Panama based implementation. So nothing changes for users who don't want to use or try it(they would continue using `PanamaVectorUtilSupport` basically). I removed all the gradle related changes and not generate any binary at all for now. To to test a user could just simply :
```
// Passes
./gradlew test
./gradlew test -Ptest.native.dotProduct=false
./gradlew test -Ptest.native.dotProduct=true -Ptests.jvmargs="-Djava.library.path=/home/simd"
./gradlew test -Ptests.jvmargs="-Dlucene.useNativeDotProduct=true -Djava.library.path=/home/simd"
// Fails
./gradlew test -Ptests.jvmargs="-Dlucene.useNativeDotProduct=true"
./gradlew test -Ptest.native.dotProduct=true
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2960708485
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
Good catch @ExE-Boss, do you want to submit a PR to fix this since you found this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3666453702
Thanks @dweiss for taking a look. I completely agree with maintaining different native bindings in Lucene can get really challenging as you mentioned and I propose we **don't ship don't ship any pre-built binaries** (or to say at least in this PR today if someone feels otherwise). I got some cross environment binaries working but this is a never ending process and we will not be able to support all permutations/combinations of the different environments.
While I got some cross-platform binaries working for few selective environments, supporting all environment permutations will be really painful. I added the build configuration to test across different environments on my end. I'm more happy to remove the Gradle build configurations for binary generation.
One important point here is that GCC's auto-vectorization of simple C code performs comparably to (and sometimes better than) hand-written NEON/SVE implementations(except some cases/runs for SVE). This allows us to keep things simple on Lucene's end and let users experiment. This is inline with what @rmuir proposed [here in this comment](https://github.com/apache/lucene/pull/13572#issuecomment-2230061569).
So here is what I propose we can do and would look like for user :
- We can maybe keep simple dot product implementation in the C code in `dotProduct.c` in the `lucene/misc` as-is but not build it.
- Instead, we let the Lucene user to optionally generate the binary on their end using `gcc -shared -O3 -march=native -funroll-loops -o <c_file>` on the platform they wish to run (or) they would need to do so for whichever compiler and environment they are targeting(we can add more documentation if req.). User could either use the default `dotProduct.c file` with auto-vectorization or get creative and provide a more performant implementation if they wish to.
- To enable the native dot product codepath they could pass the `-Dlucene.useNativeDotProduct=true` system property.
Contract: This PR would enable the interface to interact with binaries for dot product computations but providing that binary is in user's court. Users need to provide the binary with the required signature/impl:
`int32_t dot8s(int8_t vec1[], int8_t vec2[], int32_t limit);`
Let me know if this makes it clear for what it would look like from Lucene user's perspective.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2625928791
##########
build.gradle:
##########
@@ -18,6 +18,7 @@
plugins {
id "base"
id "lucene.root-project.setup"
+ id "c"
Review Comment:
this can be moved to the project which actually uses the plugin (misc?).
##########
lucene/core/build.gradle:
##########
@@ -31,6 +31,14 @@ dependencies {
moduleTestImplementation project(':lucene:test-framework')
}
+test {
+ dependsOn ':lucene:misc:buildNative'
+ systemProperty(
+ "java.library.path",
+ project(":lucene:misc").layout.buildDirectory.get().asFile.absolutePath + "/libs/dotProduct/shared"
+ )
+}
+
Review Comment:
All of these changes should be pulled into a single gradle java plugin that handles the configuration across all the involved projects. If native libs are not enabled, these changes shouldn't be applied at all. A single plugin will make it easier to see the set of changes applied globally; currently they're scattered around.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2983479755
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
Review Comment:
I opened #15866 to change this and below ones you pointed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2962527526
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
> I'd like to get rid of this magic. it slows down,.
> Ifwe don't want to do it statically
I like the suggestion, makes sense.
> I wonder why tests pass, do we really test all methods here?
We are not running it as part of the tests since we don't generate any binary (its still in the user court) and the current [c code](https://github.com/apache/lucene/blob/main/lucene/misc/src/c/dotProduct.c) only have dotProduct impl and I tested few but not all vectorutil operations. It would be good to always run the tests to ensure it catches all issues but idk if we want to go with generating a binary path. I'm thinking if its fine to not provide any platform specific binaries(abiding by what we discussed earlier) but just extend the c code to have basic VectorUtil implementations and test that on the unit test(which this PR was doing initially). Maybe you have some better ideas to how we can have better test coverage here with/without having binary? or just use fake java implementations to test each operation?
Right now this piece is bit disconnected due to no automatic test coverage running.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2962527526
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
> I'd like to get rid of this magic. it slows down,.
> Ifwe don't want to do it statically
I like the suggestion, makes sense.
> I wonder why tests pass, do we really test all methods here?
We are not running it as part of the tests since we don't generate any binary (its still in the user court) and the current [c code](https://github.com/apache/lucene/blob/main/lucene/misc/src/c/dotProduct.c) only have dotProduct impl and I tested few but not all vectorutil operations. It would be good to always run the tests to ensure it catches all issues but idk if we want to go with generating a binary path. I'm thinking if its fine to not provide any platform specific binaries(abiding by what we discussed earlier) but just extend the c code to have basic VectorUtil implementations and test that on the unit test(which this PR was doing initially). Maybe you have some better ideas to how we can have better test coverage here with/without having binary? or just use fake java implementations to test each operation?
Right now this piece is bit disconnected due to no automatic test coverage running.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3894896454
@dweiss I pushed a version that's much less intrusive and tried to simplify/refactor some bits as you suggested. Now we do not even initialize the `NativeVectorUtilSupport` class(confirmed by deliberately throwing an exception in static block) unless If the user has user has opted-in for it by passing `-Dlucene.useNativeDotProduct=true` property. So the normal/current codepath remains unchanged and continues with Panama.
Let me know your thoughts!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3901934753
Thank you @dweiss ! I'll wait for a few days to see if anyone has more feedback and then merge.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3902245260
Thank you.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3880385001
Yes, I think so.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2792293666
##########
lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java:
##########
@@ -178,9 +178,9 @@ static VectorizationProvider lookup(boolean testMode) {
}
} catch (NoSuchMethodException | IllegalAccessException e) {
throw new LinkageError(
- "PanamaVectorizationProvider is missing correctly typed constructor", e);
+ "NativeVectorizationProvider is missing correctly typed constructor", e);
} catch (ClassNotFoundException cnfe) {
- throw new LinkageError("PanamaVectorizationProvider is missing in Lucene JAR file", cnfe);
+ throw new LinkageError("NativeVectorizationProvider is missing in Lucene JAR file", cnfe);
}
Review Comment:
This will always attempt to load the native library, won't it? Maybe this native vectorization provider should be an opt-in here (controlled by an explicit property?), followed by the panama provider that we already have?
This addition really should have a minimum surface impact over the current codebase because we don't have any test support for this native wrapper.
I think it'll confuse people if we start to emit "NativeVectorizationProvider" and they don't know what it is or where to get it.
If you make it an opt-in property, you could also remove test.native.dotProduct and use that property to enable native provider in tests and benchmarks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3752311480
This PR has not had activity in the past 2 weeks, labeling it as stale. If the PR is waiting for review, notify the dev@lucene.apache.org list. Thank you for your contribution!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3873508350
FYI we now have a bug/issue filed in OpenJDK to improve vector operation support on AArch64 machines : https://bugs.openjdk.org/browse/JDK-8377386
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2793476555
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
I agree we can have some scorers(dedicated to native; something like `NativeDotProductScorer` etc) but I wonder if the current way of branching is better for a scenario where a user do not need to override all the `VectorUtil` impls and it automatically fallback to Panama if there isn't a implemented function in the library. Not having this branching would mean user needs to implement all the functions from `VectorUtil` compulsorily(since there will be no fallback). Do you think we should mandate that or maybe what you are mentioning solves that problem also or you have some ideas?.
I kinda liked the fact that I was able to override just a specific dot product implementation out of many (which is apparently slow on AArch machine) but it might not be true for all?
Let me know what you think.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2792315090
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
The same here. It adds another branch/ delegation layer... I'd be more comfortable with some code duplication (scorers) but either calling the "native" code or panama code - without all these conditionals.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "ExE-Boss (via GitHub)" <gi...@apache.org>.
ExE-Boss commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2960053681
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
Unless `mh` is wrapped [in a spreader handle][mh-asSpreader] elsewhere, this must use [`MethodHandle::invokeWithArguments`], as otherwise, this’ll throw `WrongMethodTypeException` due to the [`MethodHandle::invokeExact`] call expecting a method handle that takes exactly one argument of type `Object[]`:
```suggestion
return (T) mh.invokeWithArguments(args);
```
[mh-asSpreader]: https://docs.oracle.com/en/java/javase/25/docs/api/java.base/java/lang/invoke/MethodHandle.html#asSpreader(java.lang.Class,int)
[`MethodHandle::invokeWithArguments`]: https://docs.oracle.com/en/java/javase/25/docs/api/java.base/java/lang/invoke/MethodHandle.html#invokeWithArguments(java.lang.Object...)
[`MethodHandle::invokeExact`]: https://docs.oracle.com/en/java/javase/25/docs/api/java.base/java/lang/invoke/MethodHandle.html#invokeExact(java.lang.Object...)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "uschindler (via GitHub)" <gi...@apache.org>.
uschindler commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2962075784
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
I'd like to get rid of this magic. it slows down,. Basically you need to add a stupid try...catch on all calls due to the stupid "throws Throwable" on invokeExact, so it is a bit of copypaste but much better. we should not use dynamic invocations here.
Ifwe don't want to do it statically (and also create numeric wrappers on each call for native types), we should really better convert it to a spreader on construction on MethodHandle and then use ivokeExact().
I wonder why tests pass, do we really test all methods here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3839644363
The implementation here isn't what I had in mind @shubhamvishu - I wanted this to be a generic java service rather than a hardcoded set of options and alternatives. I'm sorry I didn't have the time to contribute here and I don't want to stand in the way - perhaps what I had in mind can be implemented as a follow-up in another patch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3675832904
It's not what I had in mind, sorry for being vague. I have a strong feeling that the entire thing should be implemented using standard ServiceLoader mechanisms. So the default singleton in VectorizationProvider.Holder.INSTANCE should be instantiated using a service and the logic of this method -
```
static VectorizationProvider lookup(boolean testMode) {
```
should be moved to individual service implementations (verification if a particular implementation can be used in the current environment, along with the instantiation of that implementation). The "singleton lookup" should only load all service providers, filter out what cannot be used for whatever reason and then pick one of the remaining candidates in their preferred order (which can be controlled by a system property, for example lucene.vectorization-provider=*,panama,default would indicate the desired ordering among available implementations).
So, by default we'd have the "default" (DefaultVectorizationProvider) fallback and "panama" (PanamaVectorizationProvider) but you could add another service implementation (native). The service would need to implement two methods - one to check if it can be used and the other to provide an instance of VectorizedProvider.
It isn't a straightforward patch because a lot of the code is currently package private and intentionally hidden (and won't allow subclassing). But I have a gut feeling it's possible and it would be a lot more elegant in the long run.
It's also related to PRs like this one - https://github.com/apache/lucene/pull/15294/ which would like to "know" which service implementation is being used. This could be the name (or class) of the service provider, for example.
I'm unfortunately away this week and won't be able to contribute directly. I would hold this patch until the above avenue can be explored though (either turns out to indeed work out nicely or won't work, for some odd reason).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2793441508
##########
lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java:
##########
@@ -178,9 +178,9 @@ static VectorizationProvider lookup(boolean testMode) {
}
} catch (NoSuchMethodException | IllegalAccessException e) {
throw new LinkageError(
- "PanamaVectorizationProvider is missing correctly typed constructor", e);
+ "NativeVectorizationProvider is missing correctly typed constructor", e);
} catch (ClassNotFoundException cnfe) {
- throw new LinkageError("PanamaVectorizationProvider is missing in Lucene JAR file", cnfe);
+ throw new LinkageError("NativeVectorizationProvider is missing in Lucene JAR file", cnfe);
}
Review Comment:
Makes sense, I'll make it opt-in here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu merged PR #15508:
URL: https://github.com/apache/lucene/pull/15508
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "uschindler (via GitHub)" <gi...@apache.org>.
uschindler commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2983205831
##########
lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java:
##########
@@ -191,6 +171,31 @@ static VectorizationProvider lookup(boolean testMode) {
return new DefaultVectorizationProvider();
}
+ static VectorizationProvider lookup(String className) {
Review Comment:
method should be private
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3668692297
I looked again and I still think a nicer way to plug this native impl. into Lucene would be to make another implementation of VectorUtilSupport ("NativeVectorUtilSupport") and then use it instead of the default Panama implementation, if it is available. If you used service lookup for picking the implementation, you could select between them when classes are initialized. Then the native implementation could go into its own module and there'd be no need for hacks like system properties, internals opened for benchmarks, etc.
I may be missing something of course, but I think it'd be a much easier way forward. What do you think?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3700601404
Thank you @dweiss for sharing your approach! This indeed is more decoupled and looks better to have separate providers convey themselves if it could be used or not based on the dynamic priority order. I've implemented your patch with an additional service implementation for the Native use case. Currently, this delegates to Panama as the fastest fallback, so the native provider is only available when Panama is usable and the required binary is present.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3837465702
Hi @dweiss , I made some more changes here and looking forward to your feedback if you get chance. Thank you!
------
Unrelated: GitHub Actions appears to be experiencing an outage - workflows are stuck in the queue. https://www.githubstatus.com/
> Update - The team continues to investigate issues causing GitHub Actions jobs on hosted runners to remain queued for extended periods, with a percentage of jobs failing. We will continue to provide updates as we make progress toward mitigation.
Feb 02, 2026 - 20:27 UTC
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2793875208
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
If that's the case then sure - leave the branching - but only in those classes that are dedicated to this "native" vector util implementation. The bottom line for me is that if it isn't available (can't be loaded) there should be no change compared to the current code - it should flow directly to panama. It should also make it easier to untagle those implementations, should someone retry my efforts with better outcome. ;)
Let me know if it makes sense.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3667003654
To expand on what exactly a Lucene user would look like on their end to test/enable this :
- **Compile** the C code (default impl or custom)
```sh
// Linux/Unix
gcc -shared -O3 -march=native -funroll-loops -o /home/simd/libdotProduct.so /home/apachelucene/lucene/misc/src/c/dotProduct.c
// macOS
gcc -shared -O3 -march=native -funroll-loops -o /home/simd/libdotProduct.dylib /home/apachelucene/lucene/misc/src/c/dotProduct.c
// Windows
gcc -shared -O3 -march=native -funroll-loops -o /home/simd/dotProduct.dll /home/apachelucene/lucene/misc/src/c/dotProduct.c
```
```c
// dotProduct.c
int32_t dot8s(int8_t vec1[], int8_t vec2[], int32_t limit) {
// dot product impl
}
```
- **Test** if everything works and dot product implementation is correct
```sh
./gradlew test // PASSES, since by default native dot product is not testes
./gradlew test -Ptest.native.dotProduct=false // PASSES, switched off
./gradlew test \
-Ptest.native.dotProduct=true \
-Ptests.jvmargs="-Djava.library.path=/home/simd" // PASSES
./gradlew test -Ptest.native.dotProduct=true // FAILS, needs library to be linked
```
- **Benchmark** against the other C implementation or Panama
```sh
java --enable-native-access=ALL-UNNAMED \
--enable-preview \
-Djava.library.path="/home/simd" \
-jar lucene/benchmark-jmh/build/benchmarks/lucene-benchmark-jmh-11.0.0-SNAPSHOT.jar \
regexp "binaryDotProductVector|dot8sNative"
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "rmuir (via GitHub)" <gi...@apache.org>.
rmuir commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3667433030
> For the record, any Java library that ships with native libs has similar integration problems. You can look at jansi or any other lib that has native components for inspiration (and to get an idea) how hairy it becomes. Here's jansi's docker-based cross-compilation build.
>
> https://github.com/fusesource/jansi/blob/master/Makefile
i'd use zig
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3663515114
Sharing some JMH benchmark results from Graviton 2 and Graviton 3 with this PR
`dot8sNative` : Native dot product implementation (with SVE/NEON/Scalar)
`dot8sNativeSimple` : Simple scalar for loop approach (letting GCC auto vectorize for the native arch)
**Graviton 2**
```
Benchmark (size) Mode Cnt Score Error Units
VectorUtilBenchmark.binaryDotProductVector 1 thrpt 15 208.089 ± 0.017 ops/us
VectorUtilBenchmark.binaryDotProductVector 128 thrpt 15 15.288 ± 0.071 ops/us
VectorUtilBenchmark.binaryDotProductVector 207 thrpt 15 9.948 ± 0.063 ops/us
VectorUtilBenchmark.binaryDotProductVector 256 thrpt 15 8.326 ± 0.030 ops/us
VectorUtilBenchmark.binaryDotProductVector 300 thrpt 15 7.063 ± 0.050 ops/us
VectorUtilBenchmark.binaryDotProductVector 512 thrpt 15 4.311 ± 0.023 ops/us
VectorUtilBenchmark.binaryDotProductVector 702 thrpt 15 3.198 ± 0.026 ops/us
VectorUtilBenchmark.binaryDotProductVector 1024 thrpt 15 2.220 ± 0.020 ops/us
VectorUtilBenchmark.dot8sNative 1 thrpt 15 86.070 ± 0.014 ops/us
VectorUtilBenchmark.dot8sNative 128 thrpt 15 84.463 ± 3.198 ops/us
VectorUtilBenchmark.dot8sNative 207 thrpt 15 49.372 ± 1.150 ops/us
VectorUtilBenchmark.dot8sNative 256 thrpt 15 70.491 ± 0.226 ops/us
VectorUtilBenchmark.dot8sNative 300 thrpt 15 44.338 ± 0.611 ops/us
VectorUtilBenchmark.dot8sNative 512 thrpt 15 43.895 ± 4.055 ops/us
VectorUtilBenchmark.dot8sNative 702 thrpt 15 27.977 ± 1.614 ops/us
VectorUtilBenchmark.dot8sNative 1024 thrpt 15 27.598 ± 0.102 ops/us
VectorUtilBenchmark.dot8sNativeSimple 1 thrpt 15 103.949 ± 0.216 ops/us
VectorUtilBenchmark.dot8sNativeSimple 128 thrpt 15 89.996 ± 3.834 ops/us
VectorUtilBenchmark.dot8sNativeSimple 207 thrpt 15 52.342 ± 0.909 ops/us
VectorUtilBenchmark.dot8sNativeSimple 256 thrpt 15 66.461 ± 3.691 ops/us
VectorUtilBenchmark.dot8sNativeSimple 300 thrpt 15 49.841 ± 1.630 ops/us
VectorUtilBenchmark.dot8sNativeSimple 512 thrpt 15 47.575 ± 0.230 ops/us
VectorUtilBenchmark.dot8sNativeSimple 702 thrpt 15 30.196 ± 2.045 ops/us
VectorUtilBenchmark.dot8sNativeSimple 1024 thrpt 15 26.870 ± 2.711 ops/us
```
**Graviton 3**
```
Benchmark (size) Mode Cnt Score Error Units
VectorUtilBenchmark.binaryDotProductVector 1 thrpt 15 449.506 ± 4.604 ops/us
VectorUtilBenchmark.binaryDotProductVector 128 thrpt 15 60.230 ± 0.007 ops/us
VectorUtilBenchmark.binaryDotProductVector 207 thrpt 15 36.598 ± 0.085 ops/us
VectorUtilBenchmark.binaryDotProductVector 256 thrpt 15 31.289 ± 0.006 ops/us
VectorUtilBenchmark.binaryDotProductVector 300 thrpt 15 26.704 ± 0.044 ops/us
VectorUtilBenchmark.binaryDotProductVector 512 thrpt 15 15.934 ± 0.001 ops/us
VectorUtilBenchmark.binaryDotProductVector 702 thrpt 15 11.607 ± 0.007 ops/us
VectorUtilBenchmark.binaryDotProductVector 1024 thrpt 15 8.041 ± 0.001 ops/us
VectorUtilBenchmark.dot8sNative 1 thrpt 15 191.014 ± 2.466 ops/us
VectorUtilBenchmark.dot8sNative 128 thrpt 15 134.566 ± 2.626 ops/us
VectorUtilBenchmark.dot8sNative 207 thrpt 15 105.161 ± 6.314 ops/us
VectorUtilBenchmark.dot8sNative 256 thrpt 15 93.163 ± 3.352 ops/us
VectorUtilBenchmark.dot8sNative 300 thrpt 15 90.764 ± 7.961 ops/us
VectorUtilBenchmark.dot8sNative 512 thrpt 15 67.553 ± 1.328 ops/us
VectorUtilBenchmark.dot8sNative 702 thrpt 15 51.275 ± 3.981 ops/us
VectorUtilBenchmark.dot8sNative 1024 thrpt 15 40.886 ± 2.880 ops/us
VectorUtilBenchmark.dot8sNativeSimple 1 thrpt 15 221.399 ± 4.357 ops/us
VectorUtilBenchmark.dot8sNativeSimple 128 thrpt 15 162.158 ± 5.077 ops/us
VectorUtilBenchmark.dot8sNativeSimple 207 thrpt 15 119.323 ± 12.108 ops/us
VectorUtilBenchmark.dot8sNativeSimple 256 thrpt 15 111.288 ± 3.256 ops/us
VectorUtilBenchmark.dot8sNativeSimple 300 thrpt 15 90.587 ± 11.066 ops/us
VectorUtilBenchmark.dot8sNativeSimple 512 thrpt 15 58.725 ± 3.419 ops/us
VectorUtilBenchmark.dot8sNativeSimple 702 thrpt 15 47.595 ± 3.692 ops/us
VectorUtilBenchmark.dot8sNativeSimple 1024 thrpt 15 36.442 ± 0.377 ops/us
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3664052353
I do have some minor comments but the large one is: how do we handle native code for end-users? Do we ship multiple precompiled binaries? Do we ship none? The current build modifications works "for you" but to compile for a matrix of all the possibilities is a bit of a nightmare.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2962527526
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
> I'd like to get rid of this magic. it slows down,.....
> :
> Ifwe don't want to do it statically ........
I like the suggestion, makes sense.
> I wonder why tests pass, do we really test all methods here?
We are not running it as part of the tests since we don't generate any binary (its still in the user court; user needs to pass the binary for tests to use it overriding whichever implementations they have) and the current [c code](https://github.com/apache/lucene/blob/main/lucene/misc/src/c/dotProduct.c) only have dotProduct impl and I tested few others but not all the existing VectorUtil methods. It would be good to always run the tests to ensure it catches all issues but idk if we want to go with generating a binary path. I'm thinking if its fine to not provide any platform specific binaries(abiding by what we discussed earlier) but just extend the c code to have basic VectorUtil implementations and test that on the unit test(which this PR was doing initially). Maybe you have some better ideas to how we can have better test coverage here with/without having binary? or just use fake java implementations to test each operation?
Right now this piece is bit disconnected due to no automatic test coverage running.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native vectorization in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2962527526
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,590 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.function.Supplier;
+import java.util.logging.Logger;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> FLOAT
+ private static final FunctionDescriptor twoPointerIntToFloat =
+ FunctionDescriptor.of(JAVA_FLOAT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle dotProductFloat$MH;
+ private static final MethodHandle squareDistanceFloat$MH;
+ private static final MethodHandle cosineFloat$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ // Each method handle with a unique native method name
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ dotProductFloat$MH = getMethodHandle("dotProductFloat", twoPointerIntToFloat);
+ squareDistanceFloat$MH = getMethodHandle("squareDistanceFloat", twoPointerIntToFloat);
+ cosineFloat$MH = getMethodHandle("cosineFloat", twoPointerIntToFloat);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ dotProductFloat$MH = null;
+ squareDistanceFloat$MH = null;
+ cosineFloat$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static float invokeFloatMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (float) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> T invokeOrDelegate(MethodHandle mh, Supplier<T> delegate, Object... args) {
+ if (mh != null) {
+ try {
+ return (T) mh.invokeExact(args);
Review Comment:
> I'd like to get rid of this magic. it slows down,.....
> :
> Ifwe don't want to do it statically ........
I like the suggestion, makes sense.
> I wonder why tests pass, do we really test all methods here?
We are not running it as part of the tests since we don't generate any binary (its still in the user court) and the current [c code](https://github.com/apache/lucene/blob/main/lucene/misc/src/c/dotProduct.c) only have dotProduct impl and I tested few but not all vectorutil operations. It would be good to always run the tests to ensure it catches all issues but idk if we want to go with generating a binary path. I'm thinking if its fine to not provide any platform specific binaries(abiding by what we discussed earlier) but just extend the c code to have basic VectorUtil implementations and test that on the unit test(which this PR was doing initially). Maybe you have some better ideas to how we can have better test coverage here with/without having binary? or just use fake java implementations to test each operation?
Right now this piece is bit disconnected due to no automatic test coverage running.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3839661971
Or you know what - please give me time until next week, I'll make another attempt then. I'm out of office this week, sorry.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3842035833
No problem at all. Thank you @dweiss for the helping hand. It would be a good addition indeed to have a generic java service. Please enjoy your time off for now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3673653537
> i'd use zig
Very Interesting! I'd love to explore this path and would be good in general if we could have some reliable way to provide efficient binaries to users. Maybe we could pursue this as a followup and keep this change just focussed on adding support to use the provided binary (next we could try adding actual binary support)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "dweiss (via GitHub)" <gi...@apache.org>.
dweiss commented on PR #15508:
URL: https://github.com/apache/lucene/pull/15508#issuecomment-3676081083
This is a quick and dirty PoC that I quickly wrote to better show what I mean -
https://github.com/apache/lucene/compare/main...dweiss:lucene:vector-prov-service?expand=1
it doesn't fully work (there are some clashes between service loader and the module system, plus a ton of cleanups to be made) but if you take a look at VectorizationProvider class it'll give you an idea what I think those losely-pluggable implementations should be loaded like.
Now... it's just an idea - I'm not really heavily advocating to switch to it but I think it'd be better in the long term if we somehow decoupled those multiple implementations (and this includes potentially removing the need for multi-release jars, which are difficult to work with).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2793476555
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
I agree we can have some scorers(dedicated to native; `NativeDotProductScorer` etc) but I wonder if this is better for a scenario where a user do not need to override all the `VectorUtil` impls and fallback to Panama if there isn't a implemented function in the library. Not having this branching would mean user needs to implement all the functions from `VectorUtil` compulsorily(since there will be no fallback). Do you think we should mandate that or maybe what you are mentioning solves that problem also or you have some ideas?.
I kinda liked the fact that I was able to override just a specific dot product implementation out of many (which is apparently slow on AArch machine) but it might not be true for all?
Let me know your thoughts.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2793476555
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
I agree we can have some scorers(dedicated to native; `NativeDotProductScorer` etc) but I wonder if this is better for a scenario where a user do not need to override all the `VectorUtil` impls and fallback to Panama if there isn't a implemented function in the library. Not having this branching would mean user needs to implement all the functions from `VectorUtil` compulsorily(since there will be no fallback). Do you think we should mandate that or maybe what you are mentioning solves that problem also or you have some ideas?.
I kinda liked the fact that I was able to override just a specific dot product implementation out of many (which is apparently slow on AArch machine) but it might not be true for all?
Let me know what you think.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2793476555
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
I agree we can have some scorers(dedicated to native; `NativeDotProductScorer` etc) but I wonder if this is better for a scenario where a user do not need to override all the VectorUtil impls and fallsback to Panama if there isn't a implemented function in the library. Not having this branching would mean user needs to implement all the functions from `VectorUtil` compulsorily(since there will be no fallback). Do you think we should mandate that or maybe what you are mentioning solves that problem also or you have some ideas?.
I kinda liked the fact that I was able to override just a specific dot product implementation out of many (which is apparently slow on AArch machine) but it might not be true for all?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org
Re: [PR] Use native dot product in Lucene [lucene]
Posted by "shubhamvishu (via GitHub)" <gi...@apache.org>.
shubhamvishu commented on code in PR #15508:
URL: https://github.com/apache/lucene/pull/15508#discussion_r2793476555
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -0,0 +1,603 @@
+/*
+ * 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.internal.vectorization;
+
+import static java.lang.foreign.ValueLayout.JAVA_BYTE;
+import static java.lang.foreign.ValueLayout.JAVA_DOUBLE;
+import static java.lang.foreign.ValueLayout.JAVA_FLOAT;
+import static java.lang.foreign.ValueLayout.JAVA_INT;
+import static java.lang.foreign.ValueLayout.JAVA_LONG;
+
+import java.lang.foreign.AddressLayout;
+import java.lang.foreign.FunctionDescriptor;
+import java.lang.foreign.Linker;
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SymbolLookup;
+import java.lang.foreign.ValueLayout;
+import java.lang.invoke.MethodHandle;
+import java.util.logging.Logger;
+import jdk.incubator.vector.FloatVector;
+import org.apache.lucene.util.Constants;
+
+/**
+ * VectorUtilSupport implementation that uses native bindings for optimized vector operations(using
+ * Foreign Function and Memory API (FFM)) if available(optional) or else fallback to
+ * PanamaVectorUtil implementations.
+ *
+ * <p>This class provides access to native C implementations of dot product operations from the
+ * loaded shared/dynamic library(.so|.dylib|.dll) which is generated from C code and linked at
+ * runtime. The native library contains multiple optimized implementations:
+ *
+ * <p>PanamaVectorUtilSupport#dotProduct use this Native C implementation for dot product
+ * calculation if system property <b>lucene.useNativeDotProduct=true</b> is passed it always tries
+ * to ensure binary is provided and required methods are implemented
+ *
+ * <p>It Uses <code>Linker.Option.critical(true)</code> for optimal performance by eliminating the
+ * overhead of ensuring MemorySegments are allocated off-heap before native calls.
+ */
+@SuppressWarnings("restricted")
+public final class NativeVectorUtilSupport implements VectorUtilSupport {
+
+ private final VectorUtilSupport delegateVectorUtilSupport;
+
+ public static final AddressLayout POINTER = ValueLayout.ADDRESS;
+
+ private static final Linker LINKER = Linker.nativeLinker();
+ private static final SymbolLookup SYMBOL_LOOKUP;
+
+ @SuppressWarnings("NonFinalStaticField")
+ private static boolean isLibraryLoaded;
+
+ // TODO: Make this dynamic?
+ public static final String NATIVE_VECTOR_LIBRARY_NAME = "dotProduct";
+
+ public NativeVectorUtilSupport(VectorUtilSupport vectorUtilSupport) {
+ this.delegateVectorUtilSupport = vectorUtilSupport;
+ }
+
+ static {
+ try {
+ // Attempt to load the library
+ System.loadLibrary(NATIVE_VECTOR_LIBRARY_NAME);
+ isLibraryLoaded = true; // If successful, set the flag to true
+ } catch (UnsatisfiedLinkError e) {
+ // If the library loading fails, set the flag to false
+ isLibraryLoaded = false;
+ Logger.getLogger(NativeVectorUtilSupport.class.getName())
+ .warning("No native library" + NATIVE_VECTOR_LIBRARY_NAME + " found : " + e.getMessage());
+ }
+ }
+
+ // Function descriptors
+ // (POINTER, POINTER, INT) -> INT
+ private static final FunctionDescriptor twoPointerIntToInt =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, INT) -> LONG
+ private static final FunctionDescriptor twoPointerIntToLong =
+ FunctionDescriptor.of(JAVA_LONG, POINTER, POINTER, JAVA_INT);
+
+ // (POINTER, POINTER, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor minMaxScalarQuantizeDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT, POINTER, POINTER, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_FLOAT, JAVA_INT);
+
+ // (POINTER, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, FLOAT, INT) -> FLOAT
+ private static final FunctionDescriptor recalculateOffsetDesc =
+ FunctionDescriptor.of(
+ JAVA_FLOAT,
+ POINTER,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_FLOAT,
+ JAVA_INT);
+
+ // (POINTER, POINTER, DOUBLE, INT) -> INT
+ private static final FunctionDescriptor filterByScoreDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, POINTER, JAVA_DOUBLE, JAVA_INT);
+
+ // (POINTER, BYTE, INT) -> POINTER
+ private static final FunctionDescriptor l2normalizeDesc =
+ FunctionDescriptor.of(POINTER, POINTER, JAVA_BYTE, JAVA_INT);
+
+ // (POINTER, INT) -> void
+ private static final FunctionDescriptor expand8Desc =
+ FunctionDescriptor.ofVoid(POINTER, JAVA_INT);
+
+ // (POINTER, INT, INT, INT) -> INT
+ private static final FunctionDescriptor findNextGEQDesc =
+ FunctionDescriptor.of(JAVA_INT, POINTER, JAVA_INT, JAVA_INT, JAVA_INT);
+
+ // Method handles
+ private static final MethodHandle dotProduct$MH;
+ private static final MethodHandle squareDistance$MH;
+ private static final MethodHandle cosine$MH;
+ private static final MethodHandle int4SquareDistance$MH;
+ private static final MethodHandle int4SquareDistanceSinglePacked$MH;
+ private static final MethodHandle int4SquareDistanceBothPacked$MH;
+ private static final MethodHandle uint8SquareDistance$MH;
+ private static final MethodHandle uint8DotProduct$MH;
+ private static final MethodHandle int4DotProduct$MH;
+ private static final MethodHandle int4DotProductSinglePacked$MH;
+ private static final MethodHandle int4DotProductBothPacked$MH;
+ private static final MethodHandle int4BitDotProduct$MH;
+ private static final MethodHandle int4DibitDotProduct$MH;
+ private static final MethodHandle minMaxScalarQuantize$MH;
+ private static final MethodHandle recalculateScalarQuantizationOffset$MH;
+ private static final MethodHandle filterByScore$MH;
+ private static final MethodHandle l2normalize$MH;
+ private static final MethodHandle expand8$MH;
+ private static final MethodHandle findNextGEQ$MH;
+
+ public static boolean isLibraryLoaded() {
+ return isLibraryLoaded;
+ }
+
+ static {
+ if (isLibraryLoaded) {
+ SymbolLookup loaderLookup = SymbolLookup.loaderLookup();
+ SYMBOL_LOOKUP = name -> loaderLookup.find(name).or(() -> LINKER.defaultLookup().find(name));
+
+ dotProduct$MH = getMethodHandle("dotProduct", twoPointerIntToInt);
+ squareDistance$MH = getMethodHandle("squareDistance", twoPointerIntToInt);
+ cosine$MH = getMethodHandle("cosine", twoPointerIntToInt);
+ int4SquareDistance$MH = getMethodHandle("int4SquareDistance", twoPointerIntToInt);
+ int4SquareDistanceSinglePacked$MH =
+ getMethodHandle("int4SquareDistanceSinglePacked", twoPointerIntToInt);
+ int4SquareDistanceBothPacked$MH =
+ getMethodHandle("int4SquareDistanceBothPacked", twoPointerIntToInt);
+ uint8SquareDistance$MH = getMethodHandle("uint8SquareDistance", twoPointerIntToInt);
+ uint8DotProduct$MH = getMethodHandle("uint8DotProduct", twoPointerIntToInt);
+ int4DotProduct$MH = getMethodHandle("int4DotProduct", twoPointerIntToInt);
+ int4DotProductSinglePacked$MH =
+ getMethodHandle("int4DotProductSinglePacked", twoPointerIntToInt);
+ int4DotProductBothPacked$MH = getMethodHandle("int4DotProductBothPacked", twoPointerIntToInt);
+ int4BitDotProduct$MH = getMethodHandle("int4BitDotProduct", twoPointerIntToLong);
+ int4DibitDotProduct$MH = getMethodHandle("int4DibitDotProduct", twoPointerIntToLong);
+ minMaxScalarQuantize$MH = getMethodHandle("minMaxScalarQuantize", minMaxScalarQuantizeDesc);
+ recalculateScalarQuantizationOffset$MH =
+ getMethodHandle("recalculateScalarQuantizationOffset", recalculateOffsetDesc);
+ filterByScore$MH = getMethodHandle("filterByScore", filterByScoreDesc);
+ l2normalize$MH = getMethodHandle("l2normalize", l2normalizeDesc);
+ expand8$MH = getMethodHandle("expand8", expand8Desc);
+ findNextGEQ$MH = getMethodHandle("findNextGEQ", findNextGEQDesc);
+ } else if (Constants.NATIVE_DOT_PRODUCT_ENABLED) {
+ throw new RuntimeException("Native library dotProduct missing!");
+ } else {
+ SYMBOL_LOOKUP = null;
+ dotProduct$MH = null;
+ squareDistance$MH = null;
+ cosine$MH = null;
+ int4SquareDistance$MH = null;
+ int4SquareDistanceSinglePacked$MH = null;
+ int4SquareDistanceBothPacked$MH = null;
+ uint8SquareDistance$MH = null;
+ uint8DotProduct$MH = null;
+ int4DotProduct$MH = null;
+ int4DotProductSinglePacked$MH = null;
+ int4DotProductBothPacked$MH = null;
+ int4BitDotProduct$MH = null;
+ int4DibitDotProduct$MH = null;
+ minMaxScalarQuantize$MH = null;
+ recalculateScalarQuantizationOffset$MH = null;
+ filterByScore$MH = null;
+ l2normalize$MH = null;
+ expand8$MH = null;
+ findNextGEQ$MH = null;
+ }
+ }
+
+ private static MethodHandle getMethodHandle(String methodName, FunctionDescriptor descriptor) {
+ MethodHandle mh =
+ SYMBOL_LOOKUP
+ .find(methodName)
+ .map(addr -> LINKER.downcallHandle(addr, descriptor, Linker.Option.critical(true)))
+ .orElse(null);
+ if (mh == null && Constants.NATIVE_STRICT_MODE) {
+ throw new RuntimeException("C code for " + methodName + " was not linked!");
+ }
+ return mh;
+ }
+
+ // Reusable invoke helpers for signatures used multiple times
+ private static int invokeIntMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (int) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ private static long invokeLongMethodHandle(MethodHandle mh, MemorySegment a, MemorySegment b) {
+ try {
+ return (long) mh.invokeExact(a, b, (int) a.byteSize());
+ } catch (Throwable ex) {
+ throw new AssertionError("should not reach here", ex);
+ }
+ }
+
+ static FloatVector fma(FloatVector a, FloatVector b, FloatVector c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ static float fma(float a, float b, float c) {
+ return PanamaVectorUtilSupport.fma(a, b, c);
+ }
+
+ public static float cosine(byte[] a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static float cosine(MemorySegment a, MemorySegment b) {
+ if (cosine$MH != null) {
+ return invokeIntMethodHandle(cosine$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.cosine(a, b);
+ }
+
+ public static int dotProduct(byte[] a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int dotProduct(MemorySegment a, MemorySegment b) {
+ if (dotProduct$MH != null) {
+ return invokeIntMethodHandle(dotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.dotProduct(a, b);
+ }
+
+ public static int squareDistance(MemorySegment a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int squareDistance(byte[] a, MemorySegment b) {
+ if (squareDistance$MH != null) {
+ return invokeIntMethodHandle(squareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.squareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(byte[] a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistance(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistance$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistance(a, b);
+ }
+
+ public static int int4SquareDistanceSinglePacked(byte[] a, MemorySegment b) {
+ if (int4SquareDistanceSinglePacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceSinglePacked$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceSinglePacked(a, b);
+ }
+
+ public static int uint8SquareDistance(byte[] a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8SquareDistance(MemorySegment a, MemorySegment b) {
+ if (uint8SquareDistance$MH != null) {
+ return invokeIntMethodHandle(uint8SquareDistance$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8SquareDistance(a, b);
+ }
+
+ public static int uint8DotProduct(byte[] a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int uint8DotProduct(MemorySegment a, MemorySegment b) {
+ if (uint8DotProduct$MH != null) {
+ return invokeIntMethodHandle(uint8DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.uint8DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(byte[] a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, MemorySegment.ofArray(a), b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProduct(MemorySegment a, MemorySegment b) {
+ if (int4DotProduct$MH != null) {
+ return invokeIntMethodHandle(int4DotProduct$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProduct(a, b);
+ }
+
+ public static int int4DotProductSinglePacked(byte[] unpacked, MemorySegment packed) {
+ if (int4DotProductSinglePacked$MH != null) {
+ return invokeIntMethodHandle(
+ int4DotProductSinglePacked$MH, MemorySegment.ofArray(unpacked), packed);
+ }
+ return PanamaVectorUtilSupport.int4DotProductSinglePacked(unpacked, packed);
+ }
+
+ public static int int4SquareDistanceBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4SquareDistanceBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4SquareDistanceBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4SquareDistanceBothPacked(a, b);
+ }
+
+ public static int int4DotProductBothPacked(MemorySegment a, MemorySegment b) {
+ if (int4DotProductBothPacked$MH != null) {
+ return invokeIntMethodHandle(int4DotProductBothPacked$MH, a, b);
+ }
+ return PanamaVectorUtilSupport.int4DotProductBothPacked(a, b);
+ }
+
+ @Override
+ public float dotProduct(float[] a, float[] b) {
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public float cosine(float[] v1, float[] v2) {
+ return delegateVectorUtilSupport.cosine(v1, v2);
+ }
+
+ @Override
+ public float squareDistance(float[] a, float[] b) {
+ return delegateVectorUtilSupport.squareDistance(a, b);
+ }
+
+ @Override
+ public int dotProduct(byte[] a, byte[] b) {
+ if (dotProduct$MH != null) {
+ return dotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.dotProduct(a, b);
+ }
+
+ @Override
+ public int int4DotProduct(byte[] a, byte[] b) {
+ if (int4DotProduct$MH != null) {
+ return int4DotProduct(MemorySegment.ofArray(a), MemorySegment.ofArray(b));
+ }
+ return delegateVectorUtilSupport.int4DotProduct(a, b);
Review Comment:
I agree we can have some scorers(dedicated to native; `NativeDotProductScorer` etc) but I wonder if this is better for a scenario where a user do not need to override all the `VectorUtil` impls and fallback to Panama if there isn't a implemented function in the library. Not having this branching would mean user needs to implement all the functions from `VectorUtil` compulsorily(since there will be no fallback). Do you think we should mandate that or maybe what you are mentioning solves that problem also or you have some ideas?.
I kinda liked the fact that I was able to override just a specific dot product implementation out of many (which is apparently slow on AArch machine) but it might not be true for all?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@lucene.apache.org
For additional commands, e-mail: issues-help@lucene.apache.org