SimdOps and NativeSimd ops refactored, VectorUtilSupport simplified (#498)

* SimdOps and NativeSimd ops refactored to remove duplicate code, VectorUtilSupport simplified for Native vs Panama

* Add back missing calls for assemble and sum and ensure assemble isn't tried with MemorySegments

* Clean up more unnessisary casts and add missing intoByteSequence call

* Use type support in SimdOps when returning a new vector

* Review points: remove SimdOps and NativeSimdOps, move logic to VectorUtilSupport

Author: Jake Luciani

jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java

@@ -225,14 +225,16 @@ public ScoreFunction.ApproximateScoreFunction precomputedScoreFunctionFor(Vector
 
     public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(VectorFloat<?> q, VectorSimilarityFunction similarityFunction) {
         VectorFloat<?> centeredQuery = pq.globalCentroid == null ? q : VectorUtil.sub(q, pq.globalCentroid);
+
+        final int subspaceCount = pq.getSubspaceCount();
         switch (similarityFunction) {
             case DOT_PRODUCT:
                 return (node2) -> {
                     var encodedChunk = getChunk(node2);
                     var encodedOffset = getOffsetInChunk(node2);
                     // compute the dot product of the query and the codebook centroids corresponding to the encoded points
                     float dp = 0;
-                    for (int m = 0; m < pq.getSubspaceCount(); m++) {
+                    for (int m = 0; m < subspaceCount; m++) {
                         int centroidIndex = Byte.toUnsignedInt(encodedChunk.get(m + encodedOffset));
                         int centroidLength = pq.subvectorSizesAndOffsets[m][0];
                         int centroidOffset = pq.subvectorSizesAndOffsets[m][1];
@@ -249,7 +251,7 @@ public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(VectorFloat<?> q,
                     // compute the dot product of the query and the codebook centroids corresponding to the encoded points
                     float sum = 0;
                     float norm2 = 0;
-                    for (int m = 0; m < pq.getSubspaceCount(); m++) {
+                    for (int m = 0; m < subspaceCount; m++) {
                         int centroidIndex = Byte.toUnsignedInt(encodedChunk.get(m + encodedOffset));
                         int centroidLength = pq.subvectorSizesAndOffsets[m][0];
                         int centroidOffset = pq.subvectorSizesAndOffsets[m][1];
@@ -267,7 +269,7 @@ public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(VectorFloat<?> q,
                     var encodedOffset = getOffsetInChunk(node2);
                     // compute the euclidean distance between the query and the codebook centroids corresponding to the encoded points
                     float sum = 0;
-                    for (int m = 0; m < pq.getSubspaceCount(); m++) {
+                    for (int m = 0; m < subspaceCount; m++) {
                         int centroidIndex = Byte.toUnsignedInt(encodedChunk.get(m + encodedOffset));
                         int centroidLength = pq.subvectorSizesAndOffsets[m][0];
                         int centroidOffset = pq.subvectorSizesAndOffsets[m][1];

jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java

@@ -115,7 +115,6 @@ public interface VectorUtilSupport {
 
   int hammingDistance(long[] v1, long[] v2);
 
-
   // default implementation used here because Panama SIMD can't express necessary SIMD operations and degrades to scalar
   /**
    * Calculates the similarity score of multiple product quantization-encoded vectors against a single query vector,

jvector-native/pom.xml

@@ -118,6 +118,11 @@
                 </exclusion>
             </exclusions>
         </dependency>
+        <dependency>
+            <groupId>io.github.jbellis</groupId>
+            <artifactId>jvector-twenty</artifactId>
+            <version>${project.version}</version>
+        </dependency>
         <dependency>
             <groupId>org.junit.jupiter</groupId>
             <artifactId>junit-jupiter-engine</artifactId>

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

@@ -16,103 +16,54 @@
 
 package io.github.jbellis.jvector.vector;
 
+import java.nio.ByteOrder;
+
 import io.github.jbellis.jvector.vector.cnative.NativeSimdOps;
 import io.github.jbellis.jvector.vector.types.ByteSequence;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
-
-import java.util.List;
+import jdk.incubator.vector.ByteVector;
+import jdk.incubator.vector.FloatVector;
+import jdk.incubator.vector.VectorMask;
+import jdk.incubator.vector.VectorSpecies;
 
 /**
  * VectorUtilSupport implementation that prefers native/Panama SIMD.
  */
-final class NativeVectorUtilSupport implements VectorUtilSupport
+final class NativeVectorUtilSupport extends PanamaVectorUtilSupport
 {
     @Override
-    public float dotProduct(VectorFloat<?> a, VectorFloat<?> b) {
-        return this.dotProduct(a, 0, b, 0, a.length());
-    }
-
-    @Override
-    public float cosine(VectorFloat<?> v1, VectorFloat<?> v2) {
-        return VectorSimdOps.cosineSimilarity((MemorySegmentVectorFloat)v1, (MemorySegmentVectorFloat)v2);
-    }
-
-    @Override
-    public float cosine(VectorFloat<?> a, int aoffset, VectorFloat<?> b, int boffset, int length) {
-        return VectorSimdOps.cosineSimilarity((MemorySegmentVectorFloat)a, aoffset, (MemorySegmentVectorFloat)b, boffset, length);
-    }
-
-    @Override
-    public float squareDistance(VectorFloat<?> a, VectorFloat<?> b) {
-         return this.squareDistance(a, 0, b, 0, a.length());
-    }
-
-    @Override
-    public float squareDistance(VectorFloat<?> a, int aoffset, VectorFloat<?> b, int boffset, int length) {
-        return VectorSimdOps.squareDistance((MemorySegmentVectorFloat) a, aoffset, (MemorySegmentVectorFloat) b, boffset, length);
-    }
-
-    @Override
-    public float dotProduct(VectorFloat<?> a, int aoffset, VectorFloat<?> b, int boffset, int length) {
-        return VectorSimdOps.dotProduct((MemorySegmentVectorFloat) a, aoffset, (MemorySegmentVectorFloat) b, boffset, length);
-    }
-
-    @Override
-    public VectorFloat<?> sum(List<VectorFloat<?>> vectors) {
-        return VectorSimdOps.sum(vectors);
-    }
-
-    @Override
-    public float sum(VectorFloat<?> vector) {
-        return VectorSimdOps.sum((MemorySegmentVectorFloat) vector);
-    }
-
-    @Override
-    public void scale(VectorFloat<?> vector, float multiplier) {
-        VectorSimdOps.scale((MemorySegmentVectorFloat) vector, multiplier);
-    }
-
-    @Override
-    public void addInPlace(VectorFloat<?> v1, VectorFloat<?> v2) {
-        VectorSimdOps.addInPlace((MemorySegmentVectorFloat)v1, (MemorySegmentVectorFloat)v2);
-    }
-
-    @Override
-    public void addInPlace(VectorFloat<?> vector, float value) {
-        VectorSimdOps.addInPlace((MemorySegmentVectorFloat) vector, value);
-    }
-
-    @Override
-    public void subInPlace(VectorFloat<?> v1, VectorFloat<?> v2) {
-        VectorSimdOps.subInPlace((MemorySegmentVectorFloat)v1, (MemorySegmentVectorFloat)v2);
+    protected FloatVector fromVectorFloat(VectorSpecies<Float> SPEC, VectorFloat<?> vector, int offset)
+    {
+        return FloatVector.fromMemorySegment(SPEC, ((MemorySegmentVectorFloat) vector).get(), vector.offset(offset), ByteOrder.LITTLE_ENDIAN);
     }
 
     @Override
-    public void subInPlace(VectorFloat<?> vector, float value) {
-        VectorSimdOps.subInPlace((MemorySegmentVectorFloat)vector, value);
+    protected FloatVector fromVectorFloat(VectorSpecies<Float> SPEC, VectorFloat<?> vector, int offset, int[] indices, int indicesOffset)
+    {
+        throw new UnsupportedOperationException("Assembly not supported with memory segments.");
     }
 
     @Override
-    public VectorFloat<?> sub(VectorFloat<?> a, VectorFloat<?> b) {
-        if (a.length() != b.length()) {
-            throw new IllegalArgumentException("Vectors must be the same length");
-        }
-        return sub(a, 0, b, 0, a.length());
+    protected void intoVectorFloat(FloatVector vector, VectorFloat<?> v, int offset)
+    {
+        vector.intoMemorySegment(((MemorySegmentVectorFloat) v).get(), v.offset(offset), ByteOrder.LITTLE_ENDIAN);
     }
 
     @Override
-    public VectorFloat<?> sub(VectorFloat<?> a, float value) {
-        return VectorSimdOps.sub((MemorySegmentVectorFloat) a, 0, value, a.length());
+    protected ByteVector fromByteSequence(VectorSpecies<Byte> SPEC, ByteSequence<?> vector, int offset)
+    {
+        return ByteVector.fromMemorySegment(SPEC, ((MemorySegmentByteSequence) vector).get(), offset, ByteOrder.LITTLE_ENDIAN);
     }
 
     @Override
-    public VectorFloat<?> sub(VectorFloat<?> a, int aOffset, VectorFloat<?> b, int bOffset, int length) {
-        return VectorSimdOps.sub((MemorySegmentVectorFloat) a, aOffset, (MemorySegmentVectorFloat) b, bOffset, length);
+    protected void intoByteSequence(ByteVector vector, ByteSequence<?> v, int offset)
+    {
+        vector.intoMemorySegment(((MemorySegmentByteSequence) v).get(), offset, ByteOrder.LITTLE_ENDIAN);
     }
 
     @Override
-    public void minInPlace(VectorFloat<?> v1, VectorFloat<?> v2) {
-        VectorSimdOps.minInPlace((MemorySegmentVectorFloat) v1, (MemorySegmentVectorFloat) v2);
+    protected void intoByteSequence(ByteVector vector, ByteSequence<?> v, int offset, VectorMask<Byte> mask) {
+        vector.intoMemorySegment(((MemorySegmentByteSequence) v).get(), offset, ByteOrder.LITTLE_ENDIAN, mask);
     }
 
     @Override
@@ -128,20 +79,6 @@ public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> b
         return NativeSimdOps.assemble_and_sum_f32_512(((MemorySegmentVectorFloat) data).get(), dataBase, ((MemorySegmentByteSequence) baseOffsets).get(), baseOffsetsOffset, baseOffsetsLength);
     }
 
-    @Override
-    public int hammingDistance(long[] v1, long[] v2) {
-        return VectorSimdOps.hammingDistance(v1, v2);
-    }
-
-    @Override
-    public float max(VectorFloat<?> vector) {
-        return VectorSimdOps.max((MemorySegmentVectorFloat) vector);
-    }
-
-    @Override
-    public float min(VectorFloat<?> vector) {
-        return VectorSimdOps.min((MemorySegmentVectorFloat) vector);
-    }
 
     @Override
     public void calculatePartialSums(VectorFloat<?> codebook, int codebookBase, int size, int clusterCount, VectorFloat<?> query, int queryOffset, VectorSimilarityFunction vsf, VectorFloat<?> partialSums) {
@@ -161,11 +98,6 @@ public void calculatePartialSums(VectorFloat<?> codebook, int codebookBase, int
         }
     }
 
-    @Override
-    public void quantizePartials(float delta, VectorFloat<?> partials, VectorFloat<?> partialBases, ByteSequence<?> quantizedPartials) {
-        VectorSimdOps.quantizePartials(delta, (MemorySegmentVectorFloat) partials, (MemorySegmentVectorFloat) partialBases, (MemorySegmentByteSequence) quantizedPartials);
-    }
-
     @Override
     public void bulkShuffleQuantizedSimilarity(ByteSequence<?> shuffles, int codebookCount, ByteSequence<?> quantizedPartials, float delta, float bestDistance, VectorSimilarityFunction vsf, VectorFloat<?> results) {
         assert shuffles.offset() == 0 : "Bulk shuffle shuffles are expected to have an offset of 0. Found: " + shuffles.offset();
@@ -199,50 +131,4 @@ 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 float nvqDotProduct8bit(VectorFloat<?> vector, ByteSequence<?> bytes, float growthRate, float midpoint, float minValue, float maxValue) {
-        return VectorSimdOps.nvqDotProduct8bit(
-                (MemorySegmentVectorFloat) vector, (MemorySegmentByteSequence) bytes,
-                growthRate, midpoint, minValue, maxValue
-        );
-    }
-
-    @Override
-    public float nvqSquareL2Distance8bit(VectorFloat<?> vector, ByteSequence<?> bytes, float growthRate, float midpoint, float minValue, float maxValue) {
-        return VectorSimdOps.nvqSquareDistance8bit(
-                (MemorySegmentVectorFloat) vector, (MemorySegmentByteSequence) bytes,
-                growthRate, midpoint, minValue, maxValue
-        );
-    }
-
-    @Override
-    public float[] nvqCosine8bit(VectorFloat<?> vector, ByteSequence<?> bytes, float growthRate, float midpoint, float minValue, float maxValue, VectorFloat<?> centroid) {
-        return VectorSimdOps.nvqCosine8bit(
-                (MemorySegmentVectorFloat) vector, (MemorySegmentByteSequence) bytes,
-                growthRate, midpoint, minValue, maxValue,
-                (MemorySegmentVectorFloat) centroid
-        );
-    }
-
-    @Override
-    public void nvqShuffleQueryInPlace8bit(VectorFloat<?> vector) {
-        VectorSimdOps.nvqShuffleQueryInPlace8bit((MemorySegmentVectorFloat) vector);
-    }
-
-    @Override
-    public void nvqQuantize8bit(VectorFloat<?> vector, float growthRate, float midpoint, float minValue, float maxValue, ByteSequence<?> destination) {
-        VectorSimdOps.nvqQuantize8bit((MemorySegmentVectorFloat) vector, growthRate, midpoint, minValue, maxValue, (MemorySegmentByteSequence) destination);
-    }
-
-    @Override
-    public float nvqLoss(VectorFloat<?> vector, float growthRate, float midpoint, float minValue, float maxValue, int nBits) {
-        return VectorSimdOps.nvqLoss((MemorySegmentVectorFloat) vector, growthRate, midpoint, minValue, maxValue, nBits);
-    }
-
-    @Override
-    public float nvqUniformLoss(VectorFloat<?> vector, float minValue, float maxValue, int nBits) {
-        return VectorSimdOps.nvqUniformLoss((MemorySegmentVectorFloat) vector, minValue, maxValue, nBits);
-    }
-
 }