[Native] BBQ Int4 to 1-bit dot product functions

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

@@ -73,6 +73,56 @@ public interface VectorSimilarityFunctions {
      */
     MethodHandle dotProductHandle7uBulkWithOffsets();
 
+    /**
+     * Produces a method handle returning the dot product of an int4 (half-byte) vector and
+     * a bit vector (one bit per element)
+     *
+     *
+     * <p> The type of the method handle will have {@code long} as return type, The type of
+     * its first and second arguments will be {@code MemorySegment}, whose contents is the
+     * vector data bytes. The third argument is the length of the vector data.
+     */
+    MethodHandle dotProductHandleI1I4();
+
+    /**
+     * Produces a method handle which computes the dot product of several vectors.
+     * This bulk operation can be used to compute the dot product between a
+     * single int4 query vector and a number of bit vectors (one bit per element),
+     *
+     * <p> The type of the method handle will have {@code void} as return type. The type of
+     * its first and second arguments will be {@code MemorySegment}, the former contains the
+     * vector data bytes for several vectors, while the latter just a single vector. The
+     * type of the third argument is an int, representing the dimensions of each vector. The
+     * type of the fourth argument is an int, representing the number of vectors in the
+     * first argument. The type of the final argument is a MemorySegment, into which the
+     * computed dot product float values will be stored.
+     */
+    MethodHandle dotProductHandleI1I4Bulk();
+
+    /**
+     * Produces a method handle which computes the dot product of several vectors.
+     * This bulk operation can be used to compute the dot product between a single int4 query
+     * vector and a subset of vectors from a dataset (array of 1-bit vectors). Each
+     * vector to include in the operation is identified by an offset inside the dataset.
+     *
+     * <p> The type of the method handle will have {@code void} as return type. The type of
+     * its arguments will be:
+     * <ol>
+     *     <li>a {@code MemorySegment} containing the vector data bytes for several vectors;
+     *     in other words, a contiguous array of vectors</li>
+     *     <li>a {@code MemorySegment} containing the vector data bytes for a single ("query") vector</li>
+     *     <li>an {@code int}, representing the dimensions of each vector</li>
+     *     <li>an {@code int}, representing the width (in bytes) of each vector. Or, in other words,
+     *     the distance in bytes between two vectors inside the first param's {@code MemorySegment}</li>
+     *     <li>a {@code MemorySegment} containing the indices of the vectors inside the first param's array
+     *     on which we'll compute the dot product</li>
+     *     <li>an {@code int}, representing the number of vectors for which we'll compute the dot product
+     *     (which is equal to the size - in number of elements - of the 5th and 7th {@code MemorySegment}s)</li>
+     *     <li>a {@code MemorySegment}, into which the computed dot product float values will be stored</li>
+     * </ol>
+     */
+    MethodHandle dotProductHandleI1I4BulkWithOffsets();
+
     /**
      * Produces a method handle returning the square distance of byte (unsigned int7) vectors.
      *

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 {
@@ -34,12 +35,19 @@ public final class JdkVectorLibrary implements VectorLibrary {
     static final MethodHandle dot7u$mh;
     static final MethodHandle dot7uBulk$mh;
     static final MethodHandle dot7uBulkWithOffsets$mh;
+
+    static final MethodHandle doti1i4$mh;
+    static final MethodHandle doti1i4Bulk$mh;
+    static final MethodHandle doti1i4BulkWithOffsets$mh;
+
     static final MethodHandle sqr7u$mh;
     static final MethodHandle sqr7uBulk$mh;
     static final MethodHandle sqr7uBulkWithOffsets$mh;
+
     static final MethodHandle dotf32$mh;
     static final MethodHandle dotf32Bulk$mh;
     static final MethodHandle dotf32BulkWithOffsets$mh;
+
     static final MethodHandle sqrf32$mh;
     static final MethodHandle sqrf32Bulk$mh;
     static final MethodHandle sqrf32BulkWithOffsets$mh;
@@ -71,15 +79,31 @@ public final class JdkVectorLibrary implements VectorLibrary {
                 dot7u$mh = downcallHandle("vec_dot7u" + suffix, intSingle, LinkerHelperUtil.critical());
                 dot7uBulk$mh = downcallHandle("vec_dot7u_bulk" + suffix, bulk, LinkerHelperUtil.critical());
                 dot7uBulkWithOffsets$mh = downcallHandle("vec_dot7u_bulk_offsets" + suffix, bulkOffsets, LinkerHelperUtil.critical());
+
+                doti1i4$mh = downcallHandle(
+                    "vec_dot_int1_int4" + suffix,
+                    FunctionDescriptor.of(JAVA_LONG, ADDRESS, ADDRESS, JAVA_INT),
+                    LinkerHelperUtil.critical()
+                );
+                doti1i4Bulk$mh = downcallHandle("vec_dot_int1_int4_bulk" + suffix, bulk, LinkerHelperUtil.critical());
+                doti1i4BulkWithOffsets$mh = downcallHandle(
+                    "vec_dot_int1_int4_bulk_offsets" + suffix,
+                    bulkOffsets,
+                    LinkerHelperUtil.critical()
+                );
+
                 sqr7u$mh = downcallHandle("vec_sqr7u" + suffix, intSingle, LinkerHelperUtil.critical());
                 sqr7uBulk$mh = downcallHandle("vec_sqr7u_bulk" + suffix, bulk, LinkerHelperUtil.critical());
                 sqr7uBulkWithOffsets$mh = downcallHandle("vec_sqr7u_bulk_offsets" + suffix, bulkOffsets, LinkerHelperUtil.critical());
+
                 dotf32$mh = downcallHandle("vec_dotf32" + suffix, floatSingle, LinkerHelperUtil.critical());
                 dotf32Bulk$mh = downcallHandle("vec_dotf32_bulk" + suffix, bulk, LinkerHelperUtil.critical());
                 dotf32BulkWithOffsets$mh = downcallHandle("vec_dotf32_bulk_offsets" + suffix, bulkOffsets, LinkerHelperUtil.critical());
+
                 sqrf32$mh = downcallHandle("vec_sqrf32" + suffix, floatSingle, LinkerHelperUtil.critical());
                 sqrf32Bulk$mh = downcallHandle("vec_sqrf32_bulk" + suffix, bulk, LinkerHelperUtil.critical());
                 sqrf32BulkWithOffsets$mh = downcallHandle("vec_sqrf32_bulk_offsets" + suffix, bulkOffsets, LinkerHelperUtil.critical());
+
                 INSTANCE = new JdkVectorSimilarityFunctions();
             } else {
                 if (caps < 0) {
@@ -90,6 +114,9 @@ public final class JdkVectorLibrary implements VectorLibrary {
                 dot7u$mh = null;
                 dot7uBulk$mh = null;
                 dot7uBulkWithOffsets$mh = null;
+                doti1i4$mh = null;
+                doti1i4Bulk$mh = null;
+                doti1i4BulkWithOffsets$mh = null;
                 sqr7u$mh = null;
                 sqr7uBulk$mh = null;
                 sqr7uBulkWithOffsets$mh = null;
@@ -148,6 +175,38 @@ static void dotProduct7uBulkWithOffsets(
             dot7uBulkWithOffsets(a, b, length, pitch, offsets, count, result);
         }
 
+        /**
+         * Computes the dot product of a given int4 vector with a give bit vector (1 bit per element).
+         *
+         * @param a      address of the bit vector
+         * @param query  address of the int4 vector
+         * @param length the vector dimensions
+         */
+        static long dotProductI1I4(MemorySegment a, MemorySegment query, int length) {
+            Objects.checkFromIndexSize(0, length * 4L, (int) query.byteSize());
+            Objects.checkFromIndexSize(0, length, (int) a.byteSize());
+            return doti1i4(a, query, length);
+        }
+
+        static void dotProductI1I4Bulk(MemorySegment dataset, MemorySegment query, int datasetVectorLengthInBytes, int count, MemorySegment result) {
+            Objects.checkFromIndexSize(0, datasetVectorLengthInBytes * count, (int) dataset.byteSize());
+            Objects.checkFromIndexSize(0, datasetVectorLengthInBytes * 4L, (int) query.byteSize());
+            Objects.checkFromIndexSize(0, count * Float.BYTES, (int) result.byteSize());
+            doti1i4Bulk(dataset, query, datasetVectorLengthInBytes, count, result);
+        }
+
+        static void dotProductI1I4BulkWithOffsets(
+            MemorySegment a,
+            MemorySegment b,
+            int length,
+            int pitch,
+            MemorySegment offsets,
+            int count,
+            MemorySegment result
+        ) {
+            doti1i4BulkWithOffsets(a, b, length, pitch, offsets, count, result);
+        }
+
         /**
          * Computes the square distance of given unsigned int7 byte vectors.
          *
@@ -288,6 +347,38 @@ private static void dot7uBulkWithOffsets(
             }
         }
 
+        private static long doti1i4(MemorySegment a, MemorySegment query, int length) {
+            try {
+                return (long) doti1i4$mh.invokeExact(a, query, length);
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
+        private static void doti1i4Bulk(MemorySegment a, MemorySegment query, int length, int count, MemorySegment result) {
+            try {
+                doti1i4Bulk$mh.invokeExact(a, query, length, count, result);
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
+        private static void doti1i4BulkWithOffsets(
+            MemorySegment a,
+            MemorySegment query,
+            int length,
+            int pitch,
+            MemorySegment offsets,
+            int count,
+            MemorySegment result
+        ) {
+            try {
+                doti1i4BulkWithOffsets$mh.invokeExact(a, query, length, pitch, offsets, count, result);
+            } catch (Throwable t) {
+                throw new AssertionError(t);
+            }
+        }
+
         private static int sqr7u(MemorySegment a, MemorySegment b, int length) {
             try {
                 return (int) sqr7u$mh.invokeExact(a, b, length);
@@ -387,6 +478,11 @@ private static void sqrf32BulkWithOffsets(
         static final MethodHandle DOT_HANDLE_7U;
         static final MethodHandle DOT_HANDLE_7U_BULK;
         static final MethodHandle DOT_HANDLE_7U_BULK_WITH_OFFSETS;
+
+        static final MethodHandle DOT_HANDLE_I1I4;
+        static final MethodHandle DOT_HANDLE_I1I4_BULK;
+        static final MethodHandle DOT_HANDLE_I1I4_BULK_WITH_OFFSETS;
+
         static final MethodHandle SQR_HANDLE_7U;
         static final MethodHandle SQR_HANDLE_7U_BULK;
         static final MethodHandle SQR_HANDLE_7U_BULK_WITH_OFFSETS;
@@ -438,6 +534,22 @@ private static void sqrf32BulkWithOffsets(
                     bulkOffsetScorer
                 );
 
+                DOT_HANDLE_I1I4 = lookup.findStatic(
+                    JdkVectorSimilarityFunctions.class,
+                    "dotProductI1I4",
+                    MethodType.methodType(long.class, MemorySegment.class, MemorySegment.class, int.class)
+                );
+                DOT_HANDLE_I1I4_BULK = lookup.findStatic(
+                    JdkVectorSimilarityFunctions.class,
+                    "dotProductI1I4Bulk",
+                    bulkScorer
+                );
+                DOT_HANDLE_I1I4_BULK_WITH_OFFSETS = lookup.findStatic(
+                    JdkVectorSimilarityFunctions.class,
+                    "dotProductI1I4BulkWithOffsets",
+                    bulkOffsetScorer
+                );
+
                 DOT_HANDLE_FLOAT32 = lookup.findStatic(JdkVectorSimilarityFunctions.class, "dotProductF32", singleFloatScorer);
                 DOT_HANDLE_FLOAT32_BULK = lookup.findStatic(JdkVectorSimilarityFunctions.class, "dotProductF32Bulk", bulkScorer);
                 DOT_HANDLE_FLOAT32_BULK_WITH_OFFSETS = lookup.findStatic(
@@ -473,6 +585,21 @@ public MethodHandle dotProductHandle7uBulkWithOffsets() {
             return DOT_HANDLE_7U_BULK_WITH_OFFSETS;
         }
 
+        @Override
+        public MethodHandle dotProductHandleI1I4() {
+            return DOT_HANDLE_I1I4;
+        }
+
+        @Override
+        public MethodHandle dotProductHandleI1I4Bulk() {
+            return DOT_HANDLE_I1I4_BULK;
+        }
+
+        @Override
+        public MethodHandle dotProductHandleI1I4BulkWithOffsets() {
+            return DOT_HANDLE_I1I4_BULK_WITH_OFFSETS;
+        }
+
         @Override
         public MethodHandle squareDistanceHandle7u() {
             return SQR_HANDLE_7U;

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

@@ -435,3 +435,83 @@ EXPORT void vec_sqrf32_bulk_offsets(
     f32_t *results) {
     sqrf32_inner_bulk<array_mapper>(a, b, dims, pitch, offsets, count, results);
 }
+
+EXPORT int64_t vec_dot_int1_int4(const int8_t* a, const int8_t* query, const int32_t length) {
+    int64_t subRet0 = 0;
+    int64_t subRet1 = 0;
+    int64_t subRet2 = 0;
+    int64_t subRet3 = 0;
+    int r = 0;
+    int upperBound = length & -sizeof(int64_t);
+    for (; r < upperBound; r += sizeof(int64_t)) {
+        int64_t value = *((int64_t*)(a + r));
+        int64_t q0 = *((int64_t*)(query + r));
+        subRet0 += __builtin_popcountll(q0 & value);
+        int64_t q1 = *((int64_t*)(query + r + length));
+        subRet1 += __builtin_popcountll(q1 & value);
+        int64_t q2 = *((int64_t*)(query + r + 2 * length));
+        subRet2 += __builtin_popcountll(q2 & value);
+        int64_t q3 = *((int64_t*)(query + r + 3 * length));
+        subRet3 += __builtin_popcountll(q3 & value);
+    }
+    upperBound = length & -sizeof(int32_t);
+    for (; r < upperBound; r += sizeof(int32_t)) {
+        int32_t value = *((int32_t*)(a + r));
+        int32_t q0 = *((int32_t*)(query + r));
+        subRet0 += __builtin_popcount(q0 & value);
+        int32_t q1 = *((int32_t*)(query + r + length));
+        subRet1 += __builtin_popcount(q1 & value);
+        int32_t q2 = *((int32_t*)(query + r + 2 * length));
+        subRet2 += __builtin_popcount(q2 & value);
+        int32_t q3 = *((int32_t*)(query + r + 3 * length));
+        subRet3 += __builtin_popcount(q3 & value);
+    }
+    for (; r < length; r++) {
+        int8_t value = *(a + r);
+        int8_t q0 = *(query + r);
+        subRet0 += __builtin_popcount(q0 & value & 0xFF);
+        int8_t q1 = *(query + r + length);
+        subRet1 += __builtin_popcount(q1 & value & 0xFF);
+        int8_t q2 = *(query + r + 2 * length);
+        subRet2 += __builtin_popcount(q2 & value & 0xFF);
+        int8_t q3 = *(query + r + 3 * length);
+        subRet3 += __builtin_popcount(q3 & value & 0xFF);
+    }
+    return subRet0 + (subRet1 << 1) + (subRet2 << 2) + (subRet3 << 3);
+}
+
+template <int64_t(*mapper)(const int32_t, const int32_t*)>
+static inline void dot_int1_int4_inner_bulk(
+    const int8_t* a,
+    const int8_t* query,
+    const int32_t length,
+    const int32_t pitch,
+    const int32_t* offsets,
+    const int32_t count,
+    f32_t* results
+) {
+    for (size_t c = 0; c < count; c++) {
+        const int8_t* a0 = a + mapper(c, offsets) * pitch;
+        results[c] = (f32_t)vec_dot_int1_int4(a0, query, length);
+    }
+}
+
+EXPORT void vec_dot_int1_int4_bulk(
+    const int8_t* a,
+    const int8_t* query,
+    const int32_t length,
+    const int32_t count,
+    f32_t* results) {
+    dot_int1_int4_inner_bulk<identity_mapper>(a, query, length, length, NULL, count, results);
+}
+
+EXPORT void vec_dot_int1_int4_bulk_offsets(
+    const int8_t* a,
+    const int8_t* query,
+    const int32_t length,
+    const int32_t pitch,
+    const int32_t* offsets,
+    const int32_t count,
+    f32_t* results) {
+    dot_int1_int4_inner_bulk<array_mapper>(a, query, length, pitch, offsets, count, results);
+}