Native OSQ scoring

benchmarks/src/main/java/org/elasticsearch/benchmark/vector/OSQScorerBenchmark.java

@@ -134,21 +134,21 @@ public void teardown() throws IOException {
         IOUtils.close(dirMmap, inMmap, dirNiofs, inNiofs);
     }
 
-    @Benchmark
-    public void scoreFromMemorySegmentOnlyVectorMmapScalar(Blackhole bh) throws IOException {
-        scoreFromMemorySegmentOnlyVector(bh, inMmap, scorerMmap);
-    }
+    // @Benchmark
+    // public void scoreFromMemorySegmentOnlyVectorMmapScalar(Blackhole bh) throws IOException {
+    // scoreFromMemorySegmentOnlyVector(bh, inMmap, scorerMmap);
+    // }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
     public void scoreFromMemorySegmentOnlyVectorMmapVect(Blackhole bh) throws IOException {
         scoreFromMemorySegmentOnlyVector(bh, inMmap, scorerMmap);
     }
 
-    @Benchmark
-    public void scoreFromMemorySegmentOnlyVectorNiofsScalar(Blackhole bh) throws IOException {
-        scoreFromMemorySegmentOnlyVector(bh, inNiofs, scorerNfios);
-    }
+    // @Benchmark
+    // public void scoreFromMemorySegmentOnlyVectorNiofsScalar(Blackhole bh) throws IOException {
+    // scoreFromMemorySegmentOnlyVector(bh, inNiofs, scorerNfios);
+    // }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
@@ -181,21 +181,21 @@ private void scoreFromMemorySegmentOnlyVector(Blackhole bh, IndexInput in, ES91O
         }
     }
 
-    @Benchmark
-    public void scoreFromMemorySegmentOnlyVectorBulkMmapScalar(Blackhole bh) throws IOException {
-        scoreFromMemorySegmentOnlyVectorBulk(bh, inMmap, scorerMmap);
-    }
+    // @Benchmark
+    // public void scoreFromMemorySegmentOnlyVectorBulkMmapScalar(Blackhole bh) throws IOException {
+    // scoreFromMemorySegmentOnlyVectorBulk(bh, inMmap, scorerMmap);
+    // }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
     public void scoreFromMemorySegmentOnlyVectorBulkMmapVect(Blackhole bh) throws IOException {
         scoreFromMemorySegmentOnlyVectorBulk(bh, inMmap, scorerMmap);
     }
 
-    @Benchmark
-    public void scoreFromMemorySegmentOnlyVectorBulkNiofsScalar(Blackhole bh) throws IOException {
-        scoreFromMemorySegmentOnlyVectorBulk(bh, inNiofs, scorerNfios);
-    }
+    // @Benchmark
+    // public void scoreFromMemorySegmentOnlyVectorBulkNiofsScalar(Blackhole bh) throws IOException {
+    // scoreFromMemorySegmentOnlyVectorBulk(bh, inNiofs, scorerNfios);
+    // }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
@@ -230,21 +230,21 @@ private void scoreFromMemorySegmentOnlyVectorBulk(Blackhole bh, IndexInput in, E
         }
     }
 
-    @Benchmark
-    public void scoreFromMemorySegmentAllBulkMmapScalar(Blackhole bh) throws IOException {
-        scoreFromMemorySegmentAllBulk(bh, inMmap, scorerMmap);
-    }
+    // @Benchmark
+    // public void scoreFromMemorySegmentAllBulkMmapScalar(Blackhole bh) throws IOException {
+    // scoreFromMemorySegmentAllBulk(bh, inMmap, scorerMmap);
+    // }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
     public void scoreFromMemorySegmentAllBulkMmapVect(Blackhole bh) throws IOException {
         scoreFromMemorySegmentAllBulk(bh, inMmap, scorerMmap);
     }
 
-    @Benchmark
-    public void scoreFromMemorySegmentAllBulkNiofsScalar(Blackhole bh) throws IOException {
-        scoreFromMemorySegmentAllBulk(bh, inNiofs, scorerNfios);
-    }
+    // @Benchmark
+    // public void scoreFromMemorySegmentAllBulkNiofsScalar(Blackhole bh) throws IOException {
+    // scoreFromMemorySegmentAllBulk(bh, inNiofs, scorerNfios);
+    // }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })

docs/changelog/134623.yaml

@@ -0,0 +1,5 @@
+pr: 134623
+summary: Native OSQ scoring
+area: Vector Search
+type: enhancement
+issues: []

libs/native/libraries/build.gradle

@@ -52,7 +52,7 @@ dependencies {
   libs "org.elasticsearch:zstd:${zstdVersion}:linux-aarch64"
   libs "org.elasticsearch:zstd:${zstdVersion}:linux-x86-64"
   libs "org.elasticsearch:zstd:${zstdVersion}:windows-x86-64"
-  libs "org.elasticsearch:vec:${vecVersion}@zip" // temporarily comment this out, if testing a locally built native lib
+  // libs "org.elasticsearch:vec:${vecVersion}@zip" // temporarily comment this out, if testing a locally built native lib
 }
 
 def extractLibs = tasks.register('extractLibs', Copy) {

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

@@ -70,4 +70,8 @@ public interface VectorSimilarityFunctions {
      * 4-byte float32 elements.
      */
     MethodHandle squareDistanceHandleFloat32();
+
+    MethodHandle int4BitDotProductHandle();
+
+    MethodHandle int4BitDotProductBulkHandle();
 }

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

@@ -25,6 +25,7 @@
 import static java.lang.foreign.ValueLayout.ADDRESS;
 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 static org.elasticsearch.nativeaccess.jdk.LinkerHelper.downcallHandle;
 
 public final class JdkVectorLibrary implements VectorLibrary {
@@ -36,6 +37,8 @@ public final class JdkVectorLibrary implements VectorLibrary {
     static final MethodHandle cosf32$mh;
     static final MethodHandle dotf32$mh;
     static final MethodHandle sqrf32$mh;
+    static final MethodHandle int4Bit$mh;
+    static final MethodHandle int4BitBulk$mh;
 
     public static final JdkVectorSimilarityFunctions INSTANCE;
 
@@ -100,6 +103,16 @@ public final class JdkVectorLibrary implements VectorLibrary {
                         LinkerHelperUtil.critical()
                     );
                 }
+                int4Bit$mh = downcallHandle(
+                    "int4Bit",
+                    FunctionDescriptor.of(JAVA_LONG, ADDRESS, ADDRESS, JAVA_LONG, JAVA_INT),
+                    LinkerHelperUtil.critical()
+                );
+                int4BitBulk$mh = downcallHandle(
+                    "int4BitBulk",
+                    FunctionDescriptor.of(JAVA_INT, ADDRESS, ADDRESS, JAVA_LONG, ADDRESS, JAVA_INT, JAVA_INT),
+                    LinkerHelperUtil.critical()
+                );
                 INSTANCE = new JdkVectorSimilarityFunctions();
             } else {
                 if (caps < 0) {
@@ -112,6 +125,8 @@ public final class JdkVectorLibrary implements VectorLibrary {
                 cosf32$mh = null;
                 dotf32$mh = null;
                 sqrf32$mh = null;
+                int4Bit$mh = null;
+                int4BitBulk$mh = null;
                 INSTANCE = null;
             }
         } catch (Throwable t) {
@@ -142,6 +157,34 @@ static int dotProduct7u(MemorySegment a, MemorySegment b, int length) {
             return dot7u(a, b, length);
         }
 
+        static long int4BitDotProd(MemorySegment a, MemorySegment b, long offset, int length) {
+            if (a.byteSize() != 4L * length) {
+                throw new IllegalArgumentException("dimensions differ: " + a.byteSize() + "!=" + 4L * length);
+            }
+            return int4Bit(a, b, offset, length);
+        }
+
+        private static long int4Bit(MemorySegment a, MemorySegment b, long offset, int length) {
+            try {
+                return (long) JdkVectorLibrary.int4Bit$mh.invokeExact(a, b, offset, length);
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
+        static int int4BitDotProdBulk(MemorySegment a, MemorySegment b, long offset, MemorySegment s, int count, int length) {
+            assert length >= 0;
+            return int4BitBulk(a, b, offset, s, count, length);
+        }
+
+        private static int int4BitBulk(MemorySegment a, MemorySegment b, long offset, MemorySegment s, int count, int length) {
+            try {
+                return (int) JdkVectorLibrary.int4BitBulk$mh.invokeExact(a, b, offset, s, count, length);
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
         /**
          * Computes the square distance of given unsigned int7 byte vectors.
          *
@@ -247,6 +290,8 @@ private static float sqrf32(MemorySegment a, MemorySegment b, int length) {
         static final MethodHandle COS_HANDLE_FLOAT32;
         static final MethodHandle DOT_HANDLE_FLOAT32;
         static final MethodHandle SQR_HANDLE_FLOAT32;
+        static final MethodHandle DOT_HANDLE_4BIT;
+        static final MethodHandle DOT_HANDLE_4BIT_BULK;
 
         static {
             try {
@@ -259,6 +304,19 @@ private static float sqrf32(MemorySegment a, MemorySegment b, int length) {
                 COS_HANDLE_FLOAT32 = lookup.findStatic(JdkVectorSimilarityFunctions.class, "cosineF32", mt);
                 DOT_HANDLE_FLOAT32 = lookup.findStatic(JdkVectorSimilarityFunctions.class, "dotProductF32", mt);
                 SQR_HANDLE_FLOAT32 = lookup.findStatic(JdkVectorSimilarityFunctions.class, "squareDistanceF32", mt);
+                mt = MethodType.methodType(long.class, MemorySegment.class, MemorySegment.class, long.class, int.class);
+                DOT_HANDLE_4BIT = lookup.findStatic(JdkVectorSimilarityFunctions.class, "int4BitDotProd", mt);
+                mt = MethodType.methodType(
+                    int.class,
+                    MemorySegment.class,
+                    MemorySegment.class,
+                    long.class,
+                    MemorySegment.class,
+                    int.class,
+                    int.class
+                );
+
+                DOT_HANDLE_4BIT_BULK = lookup.findStatic(JdkVectorSimilarityFunctions.class, "int4BitDotProdBulk", mt);
             } catch (NoSuchMethodException | IllegalAccessException e) {
                 throw new RuntimeException(e);
             }
@@ -288,5 +346,15 @@ public MethodHandle dotProductHandleFloat32() {
         public MethodHandle squareDistanceHandleFloat32() {
             return SQR_HANDLE_FLOAT32;
         }
+
+        @Override
+        public MethodHandle int4BitDotProductHandle() {
+            return DOT_HANDLE_4BIT;
+        }
+
+        @Override
+        public MethodHandle int4BitDotProductBulkHandle() {
+            return DOT_HANDLE_4BIT_BULK;
+        }
     }
 }

libs/simdvec/native/src/vec/c/aarch64/vec.c

@@ -299,3 +299,81 @@ EXPORT float sqrf32(const float *a, const float *b, size_t elementCount) {
 
     return result;
 }
+
+EXPORT int64_t int4Bit(uint8_t* query, uint8_t* doc, int64_t offset, int length) {
+   const size_t stride = (length / 8) * 8;
+   uint64_t dot_q0 = 0;
+   uint64_t dot_q1 = 0;
+   uint64_t dot_q2 = 0;
+   uint64_t dot_q3 = 0;
+   const uint8_t* doc_idx = doc + offset;
+   const uint8_t* query_j0 = query;
+   const uint8_t* query_j1 = query + length;
+   const uint8_t* query_j2 = query + 2 * length;
+   const uint8_t* query_j3 = query + 3 * length;
+   int i = 0;
+   for (; i < stride; i += 8) {
+       const uint64_t qv0 = *(const uint64_t*)(query_j0 + i);
+       const uint64_t qv1 = *(const uint64_t*)(query_j1 + i);
+       const uint64_t qv2 = *(const uint64_t*)(query_j2 + i);
+       const uint64_t qv3 = *(const uint64_t*)(query_j3 + i);
+       const uint64_t yv = *(const uint64_t*)(doc_idx + i);
+       dot_q0 += __builtin_popcountll(qv0 & yv);
+       dot_q1 += __builtin_popcountll(qv1 & yv);
+       dot_q2 += __builtin_popcountll(qv2 & yv);
+       dot_q3 += __builtin_popcountll(qv3 & yv);
+   }
+   for (; i < length; i++) {
+       const uint8_t qv0 = *(query_j0 + i);
+       const uint8_t qv1 = *(query_j1 + i);
+       const uint8_t qv2 = *(query_j2 + i);
+       const uint8_t qv3 = *(query_j3 + i);
+       const uint8_t yv = *(doc_idx + i);
+       dot_q0 += __builtin_popcountll(qv0 & yv);
+       dot_q1 += __builtin_popcountll(qv1 & yv);
+       dot_q2 += __builtin_popcountll(qv2 & yv);
+       dot_q3 += __builtin_popcountll(qv3 & yv);
+   }
+   return dot_q0 + (dot_q1 << 1) + (dot_q2 << 2) + (dot_q3 << 3);
+}
+
+EXPORT int32_t int4BitBulk(uint8_t* query, uint8_t* doc, int64_t offset, float32_t* scores, int count, int length) {
+    const size_t stride = (length / 8) * 8;
+    const uint8_t* query_j0 = query;
+    const uint8_t* query_j1 = query + length;
+    const uint8_t* query_j2 = query + 2 * length;
+    const uint8_t* query_j3 = query + 3 * length;
+    // assumption that the query bits are 4, and doc bits are singular
+    for (size_t idx = 0; idx < count; idx++) {
+      uint64_t dot_q0 = 0;
+      uint64_t dot_q1 = 0;
+      uint64_t dot_q2 = 0;
+      uint64_t dot_q3 = 0;
+      const uint8_t* doc_idx = doc + offset + idx * length;
+      int i = 0;
+      for (; i < stride; i += 8) {
+        const uint64_t qv0 = *(const uint64_t*)(query_j0 + i);
+        const uint64_t qv1 = *(const uint64_t*)(query_j1 + i);
+        const uint64_t qv2 = *(const uint64_t*)(query_j2 + i);
+        const uint64_t qv3 = *(const uint64_t*)(query_j3 + i);
+        const uint64_t yv = *(const uint64_t*)(doc_idx + i);
+        dot_q0 += __builtin_popcountll(qv0 & yv);
+        dot_q1 += __builtin_popcountll(qv1 & yv);
+        dot_q2 += __builtin_popcountll(qv2 & yv);
+        dot_q3 += __builtin_popcountll(qv3 & yv);
+      }
+      for (; i < length; i++) {
+         const uint8_t qv0 = *(query_j0 + i);
+         const uint8_t qv1 = *(query_j1 + i);
+         const uint8_t qv2 = *(query_j2 + i);
+         const uint8_t qv3 = *(query_j3 + i);
+         const uint8_t yv = *(doc_idx + i);
+         dot_q0 += __builtin_popcountll(qv0 & yv);
+         dot_q1 += __builtin_popcountll(qv1 & yv);
+         dot_q2 += __builtin_popcountll(qv2 & yv);
+         dot_q3 += __builtin_popcountll(qv3 & yv);
+      }
+      scores[idx] = (float32_t)(dot_q0 + (dot_q1 << 1) + (dot_q2 << 2) + (dot_q3 << 3));
+    }
+    return count;
+}

libs/simdvec/native/src/vec/c/amd64/vec.c

@@ -346,3 +346,11 @@ EXPORT float sqrf32(const float *a, const float *b, size_t elementCount) {
 
     return result;
 }
+
+EXPORT int64_t int4Bit(uint8_t* query, uint8_t* doc, int64_t offset, int length) {
+   return 0;
+}
+
+EXPORT int32_t int4BitBulk(uint8_t* query, uint8_t* doc, int64_t offset, float32_t* scores, size_t count, size_t dims) {
+  return 0;
+}

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

@@ -27,3 +27,7 @@ EXPORT float dotf32(const float *a, const float *b, size_t elementCount);
 
 EXPORT float sqrf32(const float *a, const float *b, size_t elementCount);
 
+EXPORT int64_t int4Bit(uint8_t* query, uint8_t* doc, int64_t offset, int length);
+
+EXPORT int32_t int4BitBulk(uint8_t* query, uint8_t* doc, int64_t offset, float32_t* scores, int count, int dims);
+

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

@@ -23,6 +23,8 @@ public class Similarities {
 
     static final MethodHandle DOT_PRODUCT_7U = DISTANCE_FUNCS.dotProductHandle7u();
     static final MethodHandle SQUARE_DISTANCE_7U = DISTANCE_FUNCS.squareDistanceHandle7u();
+    static final MethodHandle INT4_BIT_DP = DISTANCE_FUNCS.int4BitDotProductHandle();
+    static final MethodHandle INT4_BIT_DP_BULK = DISTANCE_FUNCS.int4BitDotProductBulkHandle();
 
     static int dotProduct7u(MemorySegment a, MemorySegment b, int length) {
         try {
@@ -51,4 +53,32 @@ static int squareDistance7u(MemorySegment a, MemorySegment b, int length) {
             }
         }
     }
+
+    public static long int4BitDotProduct(MemorySegment a, MemorySegment b, long offset, int length) {
+        try {
+            return (long) INT4_BIT_DP.invokeExact(a, b, offset, length);
+        } catch (Throwable e) {
+            if (e instanceof Error err) {
+                throw err;
+            } else if (e instanceof RuntimeException re) {
+                throw re;
+            } else {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    public static int int4BitDotProductBulk(MemorySegment a, MemorySegment b, long offset, MemorySegment scores, int count, int length) {
+        try {
+            return (int) INT4_BIT_DP_BULK.invokeExact(a, b, offset, scores, count, length);
+        } catch (Throwable e) {
+            if (e instanceof Error err) {
+                throw err;
+            } else if (e instanceof RuntimeException re) {
+                throw re;
+            } else {
+                throw new RuntimeException(e);
+            }
+        }
+    }
 }