Fix native implementations of PQ assembleAndSum and pqDecodedCosineSimilarity (#420)

* Fix native implementations of PQ assembleAndSum and pqDecodedCosineSimilarity

* Avoid additional slice by passing baseOffsetsOffset across native boundaries for assemble and sum functions

Author: jkni

jvector-native/src/main/c/jvector_simd.c

@@ -289,12 +289,13 @@ float euclidean_f32(int preferred_size, const float* a, int aoffset, const float
            : euclidean_f32_256(a, aoffset, b, boffset, length);
 }
 
-float assemble_and_sum_f32_512(const float* data, int dataBase, const unsigned char* baseOffsets, int baseOffsetsLength) {
+float assemble_and_sum_f32_512(const float* data, int dataBase, const unsigned char* baseOffsets, int baseOffsetsOffset, int baseOffsetsLength) {
     __m512 sum = _mm512_setzero_ps();
     int i = 0;
     int limit = baseOffsetsLength - (baseOffsetsLength % 16);
     __m512i indexRegister = initialIndexRegister;
     __m512i dataBaseVec = _mm512_set1_epi32(dataBase);
+    baseOffsets = baseOffsets + baseOffsetsOffset;
 
     for (; i < limit; i += 16) {
         __m128i baseOffsetsRaw = _mm_loadu_si128((__m128i *)(baseOffsets + i));
@@ -319,13 +320,14 @@ float assemble_and_sum_f32_512(const float* data, int dataBase, const unsigned c
     return res;
 }
 
-float pq_decoded_cosine_similarity_f32_512(const unsigned char* baseOffsets, int baseOffsetsLength, int clusterCount, const float* partialSums, const float* aMagnitude, float bMagnitude) {
+float pq_decoded_cosine_similarity_f32_512(const unsigned char* baseOffsets, int baseOffsetsOffset, int baseOffsetsLength, int clusterCount, const float* partialSums, const float* aMagnitude, float bMagnitude) {
     __m512 sum = _mm512_setzero_ps();
     __m512 vaMagnitude = _mm512_setzero_ps();
     int i = 0;
     int limit = baseOffsetsLength - (baseOffsetsLength % 16);
     __m512i indexRegister = initialIndexRegister;
     __m512i scale = _mm512_set1_epi32(clusterCount);
+    baseOffsets = baseOffsets + baseOffsetsOffset;
 
 
     for (; i < limit; i += 16) {

jvector-native/src/main/c/jvector_simd.h

@@ -28,8 +28,8 @@ float euclidean_f32(int preferred_size, const float* a, int aoffset, const float
 void bulk_quantized_shuffle_dot_f32_512(const unsigned char* shuffles, int codebookCount, const char* quantizedPartials, float delta, float minDistance, float* results);
 void bulk_quantized_shuffle_euclidean_f32_512(const unsigned char* shuffles, int codebookCount, const char* quantizedPartials, float delta, float minDistance, float* results);
 void bulk_quantized_shuffle_cosine_f32_512(const unsigned char* shuffles, int codebookCount, const char* quantizedPartialSums, float sumDelta, float minDistance, const char* quantizedPartialMagnitudes, float magnitudeDelta, float minMagnitude, float queryMagnitudeSquared, float* results);
-float assemble_and_sum_f32_512(const float* data, int dataBase, const unsigned char* baseOffsets, int baseOffsetsLength);
-float pq_decoded_cosine_similarity_f32_512(const unsigned char* baseOffsets, int baseOffsetsLength, int clusterCount, const float* partialSums, const float* aMagnitude, float bMagnitude);
+float assemble_and_sum_f32_512(const float* data, int dataBase, const unsigned char* baseOffsets, int baseOffsetsOffset, int baseOffsetsLength);
+float pq_decoded_cosine_similarity_f32_512(const unsigned char* baseOffsets, int baseOffsetsOffset, int baseOffsetsLength, int clusterCount, const float* partialSums, const float* aMagnitude, float bMagnitude);
 void calculate_partial_sums_dot_f32_512(const float* codebook, int codebookBase, int size, int clusterCount, const float* query, int queryOffset, float* partialSums);
 void calculate_partial_sums_euclidean_f32_512(const float* codebook, int codebookBase, int size, int clusterCount, const float* query, int queryOffset, float* partialSums);
 void calculate_partial_sums_best_dot_f32_512(const float* codebook, int codebookBase, int size, int clusterCount, const float* query, int queryOffset, float* partialSums, float* partialBestDistances);

jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java

@@ -117,16 +117,15 @@ public void minInPlace(VectorFloat<?> v1, VectorFloat<?> v2) {
 
     @Override
     public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> baseOffsets) {
-        assert baseOffsets.offset() == 0 : "Base offsets are expected to have an offset of 0. Found: " + baseOffsets.offset();
-        return NativeSimdOps.assemble_and_sum_f32_512(((MemorySegmentVectorFloat)data).get(), dataBase, ((MemorySegmentByteSequence)baseOffsets).get(), baseOffsets.length());
+        return assembleAndSum(data, dataBase, baseOffsets, 0, baseOffsets.length());
     }
 
     @Override
     public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> baseOffsets, int baseOffsetsOffset, int baseOffsetsLength)
     {
-        assert baseOffsetsOffset == 0;
-        assert baseOffsetsLength == baseOffsets.length();
-        return assembleAndSum(data, dataBase, baseOffsets);
+        assert baseOffsets.offset() == 0 : "Base offsets are expected to have an offset of 0. Found: " + baseOffsets.offset();
+        // baseOffsets is a pointer into a PQ chunk - we need to index into it by baseOffsetsOffset and provide baseOffsetsLength to the native code
+        return NativeSimdOps.assemble_and_sum_f32_512(((MemorySegmentVectorFloat) data).get(), dataBase, ((MemorySegmentByteSequence) baseOffsets).get(), baseOffsetsOffset, baseOffsetsLength);
     }
 
     @Override
@@ -189,9 +188,16 @@ public void bulkShuffleQuantizedSimilarityCosine(ByteSequence<?> shuffles, int c
 
     @Override
     public float pqDecodedCosineSimilarity(ByteSequence<?> encoded, int clusterCount, VectorFloat<?> partialSums, VectorFloat<?> aMagnitude, float bMagnitude)
+    {
+        return pqDecodedCosineSimilarity(encoded, 0, encoded.length(), clusterCount, partialSums, aMagnitude, bMagnitude);
+    }
+
+    @Override
+    public float pqDecodedCosineSimilarity(ByteSequence<?> encoded, int encodedOffset, int encodedLength, int clusterCount, VectorFloat<?> partialSums, VectorFloat<?> aMagnitude, float bMagnitude)
     {
         assert encoded.offset() == 0 : "Bulk shuffle shuffles are expected to have an offset of 0. Found: " + encoded.offset();
-        return NativeSimdOps.pq_decoded_cosine_similarity_f32_512(((MemorySegmentByteSequence) encoded).get(), encoded.length(), clusterCount, ((MemorySegmentVectorFloat) partialSums).get(), ((MemorySegmentVectorFloat) aMagnitude).get(), bMagnitude);
+        // encoded is a pointer into a PQ chunk - we need to index into it by encodedOffset and provide encodedLength to the native code
+        return NativeSimdOps.pq_decoded_cosine_similarity_f32_512(((MemorySegmentByteSequence) encoded).get(), encodedOffset, encodedLength, clusterCount, ((MemorySegmentVectorFloat) partialSums).get(), ((MemorySegmentVectorFloat) aMagnitude).get(), bMagnitude);
     }
 
     @Override

jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java

@@ -408,6 +408,7 @@ private static class assemble_and_sum_f32_512 {
             NativeSimdOps.C_POINTER,
             NativeSimdOps.C_INT,
             NativeSimdOps.C_POINTER,
+            NativeSimdOps.C_INT,
             NativeSimdOps.C_INT
         );
 
@@ -419,7 +420,7 @@ private static class assemble_and_sum_f32_512 {
     /**
      * Function descriptor for:
      * {@snippet lang=c :
-     * float assemble_and_sum_f32_512(const float *data, int dataBase, const unsigned char *baseOffsets, int baseOffsetsLength)
+     * float assemble_and_sum_f32_512(const float *data, int dataBase, const unsigned char *baseOffsets, int baseOffsetsOffset, int baseOffsetsLength)
      * }
      */
     public static FunctionDescriptor assemble_and_sum_f32_512$descriptor() {
@@ -429,24 +430,24 @@ private static class assemble_and_sum_f32_512 {
     /**
      * Downcall method handle for:
      * {@snippet lang=c :
-     * float assemble_and_sum_f32_512(const float *data, int dataBase, const unsigned char *baseOffsets, int baseOffsetsLength)
+     * float assemble_and_sum_f32_512(const float *data, int dataBase, const unsigned char *baseOffsets, int baseOffsetsOffset, int baseOffsetsLength)
      * }
      */
     public static MethodHandle assemble_and_sum_f32_512$handle() {
         return assemble_and_sum_f32_512.HANDLE;
     }
     /**
      * {@snippet lang=c :
-     * float assemble_and_sum_f32_512(const float *data, int dataBase, const unsigned char *baseOffsets, int baseOffsetsLength)
+     * float assemble_and_sum_f32_512(const float *data, int dataBase, const unsigned char *baseOffsets, int baseOffsetsOffset, int baseOffsetsLength)
      * }
      */
-    public static float assemble_and_sum_f32_512(MemorySegment data, int dataBase, MemorySegment baseOffsets, int baseOffsetsLength) {
+    public static float assemble_and_sum_f32_512(MemorySegment data, int dataBase, MemorySegment baseOffsets, int baseOffsetsOffset, int baseOffsetsLength) {
         var mh$ = assemble_and_sum_f32_512.HANDLE;
         try {
             if (TRACE_DOWNCALLS) {
-                traceDowncall("assemble_and_sum_f32_512", data, dataBase, baseOffsets, baseOffsetsLength);
+                traceDowncall("assemble_and_sum_f32_512", data, dataBase, baseOffsets, baseOffsetsOffset, baseOffsetsLength);
             }
-            return (float)mh$.invokeExact(data, dataBase, baseOffsets, baseOffsetsLength);
+            return (float)mh$.invokeExact(data, dataBase, baseOffsets, baseOffsetsOffset, baseOffsetsLength);
         } catch (Throwable ex$) {
            throw new AssertionError("should not reach here", ex$);
         }
@@ -458,6 +459,7 @@ private static class pq_decoded_cosine_similarity_f32_512 {
             NativeSimdOps.C_POINTER,
             NativeSimdOps.C_INT,
             NativeSimdOps.C_INT,
+            NativeSimdOps.C_INT,
             NativeSimdOps.C_POINTER,
             NativeSimdOps.C_POINTER,
             NativeSimdOps.C_FLOAT
@@ -471,7 +473,7 @@ private static class pq_decoded_cosine_similarity_f32_512 {
     /**
      * Function descriptor for:
      * {@snippet lang=c :
-     * float pq_decoded_cosine_similarity_f32_512(const unsigned char *baseOffsets, int baseOffsetsLength, int clusterCount, const float *partialSums, const float *aMagnitude, float bMagnitude)
+     * float pq_decoded_cosine_similarity_f32_512(const unsigned char *baseOffsets, int baseOffsetsOffset, int baseOffsetsLength, int clusterCount, const float *partialSums, const float *aMagnitude, float bMagnitude)
      * }
      */
     public static FunctionDescriptor pq_decoded_cosine_similarity_f32_512$descriptor() {
@@ -481,24 +483,24 @@ private static class pq_decoded_cosine_similarity_f32_512 {
     /**
      * Downcall method handle for:
      * {@snippet lang=c :
-     * float pq_decoded_cosine_similarity_f32_512(const unsigned char *baseOffsets, int baseOffsetsLength, int clusterCount, const float *partialSums, const float *aMagnitude, float bMagnitude)
+     * float pq_decoded_cosine_similarity_f32_512(const unsigned char *baseOffsets, int baseOffsetsOffset, int baseOffsetsLength, int clusterCount, const float *partialSums, const float *aMagnitude, float bMagnitude)
      * }
      */
     public static MethodHandle pq_decoded_cosine_similarity_f32_512$handle() {
         return pq_decoded_cosine_similarity_f32_512.HANDLE;
     }
     /**
      * {@snippet lang=c :
-     * float pq_decoded_cosine_similarity_f32_512(const unsigned char *baseOffsets, int baseOffsetsLength, int clusterCount, const float *partialSums, const float *aMagnitude, float bMagnitude)
+     * float pq_decoded_cosine_similarity_f32_512(const unsigned char *baseOffsets, int baseOffsetsOffset, int baseOffsetsLength, int clusterCount, const float *partialSums, const float *aMagnitude, float bMagnitude)
      * }
      */
-    public static float pq_decoded_cosine_similarity_f32_512(MemorySegment baseOffsets, int baseOffsetsLength, int clusterCount, MemorySegment partialSums, MemorySegment aMagnitude, float bMagnitude) {
+    public static float pq_decoded_cosine_similarity_f32_512(MemorySegment baseOffsets, int baseOffsetsOffset, int baseOffsetsLength, int clusterCount, MemorySegment partialSums, MemorySegment aMagnitude, float bMagnitude) {
         var mh$ = pq_decoded_cosine_similarity_f32_512.HANDLE;
         try {
             if (TRACE_DOWNCALLS) {
-                traceDowncall("pq_decoded_cosine_similarity_f32_512", baseOffsets, baseOffsetsLength, clusterCount, partialSums, aMagnitude, bMagnitude);
+                traceDowncall("pq_decoded_cosine_similarity_f32_512", baseOffsets, baseOffsetsOffset, baseOffsetsLength, clusterCount, partialSums, aMagnitude, bMagnitude);
             }
-            return (float)mh$.invokeExact(baseOffsets, baseOffsetsLength, clusterCount, partialSums, aMagnitude, bMagnitude);
+            return (float)mh$.invokeExact(baseOffsets, baseOffsetsOffset, baseOffsetsLength, clusterCount, partialSums, aMagnitude, bMagnitude);
         } catch (Throwable ex$) {
            throw new AssertionError("should not reach here", ex$);
         }