Wire calculatePartialSums to native SIMD via Panama FFI downcall

Author: Raghuveer Devulapalli

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

@@ -104,4 +104,17 @@ public float pqDecodedCosineSimilarity(ByteSequence<?> encoded, int encodedOffse
         // 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
+    public void calculatePartialSums(VectorFloat<?> codebook, int codebookIndex, int size, int clusterCount, VectorFloat<?> query, int queryOffset, VectorSimilarityFunction vsf, VectorFloat<?> partialSums) {
+        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;
+            default -> throw new UnsupportedOperationException("Unsupported similarity function " + vsf);
+        };
+        NativeSimdOps.calculate_partial_sums_f32_512(nativeCodebook, codebookIndex, size, clusterCount, nativeQuery, queryOffset, similarityFunction, nativePartialSums);
+    }
 }

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

@@ -847,4 +847,36 @@ public static void calculate_partial_sums_best_euclidean_f32_512(MemorySegment c
             throw new AssertionError("should not reach here", ex$);
         }
     }
+
+    private static class calculate_partial_sums_f32_512 {
+        public static final FunctionDescriptor DESC = FunctionDescriptor.ofVoid(
+                NativeSimdOps.C_POINTER,
+                NativeSimdOps.C_INT,
+                NativeSimdOps.C_INT,
+                NativeSimdOps.C_INT,
+                NativeSimdOps.C_POINTER,
+                NativeSimdOps.C_INT,
+                NativeSimdOps.C_INT,
+                NativeSimdOps.C_POINTER
+        );
+        public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("calculate_partial_sums_f32_512");
+        public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true));
+    }
+
+    /**
+     * {@snippet lang=c :
+     * void calculate_partial_sums_f32_512(const float *codebook, int codebookIndex, int size, int clusterCount, const float *query, int queryOffset, int similarityFunction, float *partialSums)
+     * }
+     */
+    public static void calculate_partial_sums_f32_512(MemorySegment codebook, int codebookIndex, int size, int clusterCount, MemorySegment query, int queryOffset, int similarityFunction, MemorySegment partialSums) {
+        var mh$ = calculate_partial_sums_f32_512.HANDLE;
+        try {
+            if (TRACE_DOWNCALLS) {
+                traceDowncall("calculate_partial_sums_f32_512", codebook, codebookIndex, size, clusterCount, query, queryOffset, similarityFunction, partialSums);
+            }
+            mh$.invokeExact(codebook, codebookIndex, size, clusterCount, query, queryOffset, similarityFunction, partialSums);
+        } catch (Throwable ex$) {
+            throw new AssertionError("should not reach here", ex$);
+        }
+    }
 }