Remove calculate_partial_sums_f32_512 dispatcher, expose variants directly

Author: r-devulap

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

@@ -259,7 +259,7 @@ JV_INLINE float euclidean_f32(const float* a, int aoffset, const float* b, int b
            : euclidean_f32_256(a, aoffset, b, boffset, length);
 }
 
-JV_INLINE void calculate_partial_sums_dot_f32_512(const float* codebook, int codebookIndex, int size, int clusterCount, const float* query, int queryOffset, float* partialSums) {
+void calculate_partial_sums_dot_f32_512(const float* codebook, int codebookIndex, int size, int clusterCount, const float* query, int queryOffset, float* partialSums) {
     int codebookBase = codebookIndex * clusterCount;
     float tempdat[16];
     if (size == 2) {
@@ -365,7 +365,7 @@ JV_INLINE void calculate_partial_sums_dot_f32_512(const float* codebook, int cod
     }
 }
 
-JV_INLINE void calculate_partial_sums_euclidean_f32_512(const float* codebook, int codebookIndex, int size, int clusterCount, const float* query, int queryOffset, float* partialSums) {
+void calculate_partial_sums_euclidean_f32_512(const float* codebook, int codebookIndex, int size, int clusterCount, const float* query, int queryOffset, float* partialSums) {
     int codebookBase = codebookIndex * clusterCount;
     float tempdat[16];
     if (size == 2) {
@@ -728,15 +728,4 @@ float pq_decoded_cosine_similarity_f32_512(const unsigned char* baseOffsets, int
     return sumResult / sqrtf(aMagnitudeResult * bMagnitude);
 }
 
-void calculate_partial_sums_f32_512(const float* codebook, int codebookIndex, int size, int clusterCount, const float* query, int queryOffset, int similarityFunction, float* partialSums) {
-    switch (similarityFunction) {
-        case 0:
-            calculate_partial_sums_euclidean_f32_512(codebook, codebookIndex, size, clusterCount, query, queryOffset, partialSums);
-            break;
-        case 1:
-            calculate_partial_sums_dot_f32_512(codebook, codebookIndex, size, clusterCount, query, queryOffset, partialSums);
-            break;
-        default:
-            break;
-    }
-}
+

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

@@ -27,5 +27,6 @@ bool check_avx512_compatibility(void);
 // APIs exposed to Java via FFI
 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_f32_512(const float* codebook, int codebookBase, int size, int clusterCount, const float* query, int queryOffset, int similarityFunction, 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_dot_f32_512(const float* codebook, int codebookBase, int size, int clusterCount, const float* query, int queryOffset, float* partialSums);
 #endif

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

@@ -110,11 +110,10 @@ public void calculatePartialSums(VectorFloat<?> codebook, int codebookIndex, int
         var nativeCodebook = ((MemorySegmentVectorFloat) codebook).get();
         var nativeQuery = ((MemorySegmentVectorFloat) query).get();
         var nativePartialSums = ((MemorySegmentVectorFloat) partialSums).get();
-        int similarityFunction = switch (vsf) {
-            case EUCLIDEAN -> 0;
-            case DOT_PRODUCT -> 1;
+        switch (vsf) {
+            case EUCLIDEAN -> NativeSimdOps.calculate_partial_sums_euclidean_f32_512(nativeCodebook, codebookIndex, size, clusterCount, nativeQuery, queryOffset, nativePartialSums);
+            case DOT_PRODUCT -> NativeSimdOps.calculate_partial_sums_dot_f32_512(nativeCodebook, codebookIndex, size, clusterCount, nativeQuery, queryOffset, nativePartialSums);
             default -> throw new UnsupportedOperationException("Unsupported similarity function " + vsf);
-        };
-        NativeSimdOps.calculate_partial_sums_f32_512(nativeCodebook, codebookIndex, size, clusterCount, nativeQuery, queryOffset, similarityFunction, nativePartialSums);
+        }
     }
 }