From 7d8c448ea8753f8e0722eacc8fbe0390a031b9cc Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Thu, 16 Jul 2026 10:18:24 +0000 Subject: [PATCH 1/9] Add RandomAccessByteVectorValues interface and ListRandomAccessByteVectorValues MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a byte-vector parallel to RandomAccessVectorValues so the HNSW pipeline can operate natively on int8 (ByteSequence) vectors without any float32 round-trip. Why new classes instead of overloading ListRandomAccessVectorValues: - RandomAccessVectorValues.getVector() is contractually fixed to return VectorFloat. Adding a constructor that accepts List> would leave getVector() unable to return the stored bytes without either a ClassCastException at runtime or silent dequantization on every access — both defeating the purpose of a native int8 path. - ByteSequence and VectorFloat are unrelated types with no common parent, so Java's type system offers no return-type covariance that could make a single getVector() work for both. - A separate interface (RandomAccessByteVectorValues) keeps the byte-vector path type-safe end-to-end and ensures that no existing RandomAccessVectorValues consumer can accidentally receive a ByteSequence where it expects a VectorFloat. New files: - RandomAccessByteVectorValues: interface mirroring RandomAccessVectorValues but with getVector() -> ByteSequence; includes the same threadLocalSupplier() default with shared/non-shared logic (shared=false => returns this, shared=true => wraps in ExplicitThreadLocal). - ListRandomAccessByteVectorValues: List>-backed implementation; isValueShared()==false, copy() returns this. --- .../ListRandomAccessByteVectorValues.java | 70 ++++++++++++++++++ .../graph/RandomAccessByteVectorValues.java | 74 +++++++++++++++++++ 2 files changed, 144 insertions(+) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java new file mode 100644 index 000000000..134b13546 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ListRandomAccessByteVectorValues.java @@ -0,0 +1,70 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.vector.types.ByteSequence; + +import java.util.List; + +/** + * A List-backed implementation of the {@link RandomAccessByteVectorValues} interface. + *

+ * It is acceptable to provide this class to a GraphBuilder, and then continue + * to add vectors to the backing List as you add to the graph. + *

+ * This will be as threadsafe as the provided List. + */ +public class ListRandomAccessByteVectorValues implements RandomAccessByteVectorValues { + private final List> vectors; + private final int dimension; + + /** + * Construct a new instance of {@link ListRandomAccessByteVectorValues}. + * + * @param vectors a (potentially mutable) list of byte vectors. + * @param dimension the dimension of the vectors. + */ + public ListRandomAccessByteVectorValues(List> vectors, int dimension) { + this.vectors = vectors; + this.dimension = dimension; + } + + @Override + public int size() { + return vectors.size(); + } + + @Override + public int dimension() { + return dimension; + } + + @Override + public ByteSequence getVector(int nodeId) { + return vectors.get(nodeId); + } + + @Override + public boolean isValueShared() { + return false; + } + + @Override + public ListRandomAccessByteVectorValues copy() { + return this; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java new file mode 100644 index 000000000..543ed47a2 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessByteVectorValues.java @@ -0,0 +1,74 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import io.github.jbellis.jvector.util.ExplicitThreadLocal; +import io.github.jbellis.jvector.vector.types.ByteSequence; + +import java.util.function.Supplier; +import java.util.logging.Logger; + +/** + * Provides random access to byte (int8) vectors by dense ordinal. + *

+ * This is the byte-vector parallel to {@link RandomAccessVectorValues}. + * It is used by graph-based index builders and searchers that operate natively + * on int8 vectors without a float32 round-trip. + */ +public interface RandomAccessByteVectorValues { + Logger LOG = Logger.getLogger(RandomAccessByteVectorValues.class.getName()); + + /** Return the number of vector values. */ + int size(); + + /** Return the dimension of the returned vector values. */ + int dimension(); + + /** + * Return the byte vector indexed at the given ordinal. + * + * @param nodeId a valid ordinal, ≥ 0 and < {@link #size()}. + */ + ByteSequence getVector(int nodeId); + + /** + * @return true iff the vector returned by {@link #getVector} is shared across calls. + * A shared vector is only valid until the next call to {@link #getVector} overwrites it. + */ + boolean isValueShared(); + + /** + * Creates a new copy of this {@link RandomAccessByteVectorValues}. + * Un-shared implementations may simply return {@code this}. + */ + RandomAccessByteVectorValues copy(); + + /** + * Returns a supplier of thread-local copies of the RABVV. + */ + default Supplier threadLocalSupplier() { + if (!isValueShared()) { + return () -> this; + } + + if (this instanceof AutoCloseable) { + LOG.warning("RABVV is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks"); + } + var tl = ExplicitThreadLocal.withInitial(this::copy); + return tl::get; + } +} From 088301505d2e86eb6960174c95bf9d9dac93e127 Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Thu, 16 Jul 2026 10:18:41 +0000 Subject: [PATCH 2/9] Add byte-similarity methods to VectorUtilSupport / VectorUtil / DefaultVectorUtilSupport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add three fundamental signed int8 vector similarity operations to the vectorization support layer so they participate in the same provider dispatch as float similarity. Bytes are treated as signed int8 (Java byte range -128..127). VectorUtilSupport — three new abstract methods: float dotProduct(ByteSequence a, ByteSequence b) float squareDistance(ByteSequence a, ByteSequence b) float cosine(ByteSequence a, ByteSequence b) VectorUtil — three new public static delegates: dotProduct(ByteSequence, ByteSequence) -> impl.dotProduct squareL2Distance(ByteSequence, ByteSequence) -> impl.squareDistance cosine(ByteSequence, ByteSequence) -> impl.cosine DefaultVectorUtilSupport — scalar loop implementations: dotProduct: accumulate (int)a.get(i) * (int)b.get(i), return as float. squareDistance: accumulate (diff * diff) for each signed byte difference. cosine: dot / sqrt(normA * normB) using per-element float promotion. PanamaVectorUtilSupport — scalar stub overrides identical to Default, so jvector-twenty compiles without requiring a SIMD implementation now. SIMD optimisation of byte similarity is a future concern. --- .../vector/DefaultVectorUtilSupport.java | 31 +++++++++++++++++++ .../jbellis/jvector/vector/VectorUtil.java | 15 +++++++++ .../jvector/vector/VectorUtilSupport.java | 9 ++++++ .../vector/PanamaVectorUtilSupport.java | 31 +++++++++++++++++++ 4 files changed, 86 insertions(+) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java index 5843dc5f6..e5f7b0824 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java @@ -338,6 +338,37 @@ public float assembleAndSumPQ( return res; } + @Override + public float dotProduct(ByteSequence a, ByteSequence b) { + float sum = 0; + for (int i = 0; i < a.length(); i++) { + sum += (int) a.get(i) * (int) b.get(i); + } + return sum; + } + + @Override + public float squareDistance(ByteSequence a, ByteSequence b) { + float sum = 0; + for (int i = 0; i < a.length(); i++) { + float diff = a.get(i) - b.get(i); + sum += diff * diff; + } + return sum; + } + + @Override + public float cosine(ByteSequence a, ByteSequence b) { + float dot = 0, normA = 0, normB = 0; + for (int i = 0; i < a.length(); i++) { + float ai = a.get(i), bi = b.get(i); + dot += ai * bi; + normA += ai * ai; + normB += bi * bi; + } + return (float) (dot / Math.sqrt(normA * normB)); + } + @Override public int hammingDistance(long[] v1, long[] v2) { int hd = 0; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java index 744d5ec75..01550f264 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java @@ -174,6 +174,21 @@ public static float assembleAndSumPQ(VectorFloat data, int subspaceCount, Byt return impl.assembleAndSumPQ(data, subspaceCount, dataOffsets1, dataOffsetsOffset1, dataOffsets2, dataOffsetsOffset2, clusterCount); } + /** Returns the dot product of two signed int8 byte vectors. */ + public static float dotProduct(ByteSequence a, ByteSequence b) { + return impl.dotProduct(a, b); + } + + /** Returns the sum of squared differences of two signed int8 byte vectors. */ + public static float squareL2Distance(ByteSequence a, ByteSequence b) { + return impl.squareDistance(a, b); + } + + /** Returns the cosine similarity of two signed int8 byte vectors. */ + public static float cosine(ByteSequence a, ByteSequence b) { + return impl.cosine(a, b); + } + public static int hammingDistance(long[] v1, long[] v2) { return impl.hammingDistance(v1, v2); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java index 118f16ca6..01a706405 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java @@ -130,6 +130,15 @@ public interface VectorUtilSupport { */ float assembleAndSumPQ(VectorFloat codebookPartialSums, int subspaceCount, ByteSequence vector1Ordinals, int vector1OrdinalOffset, ByteSequence node2Ordinals, int node2OrdinalOffset, int clusterCount); + /** Calculates the dot product of two signed int8 byte vectors. */ + float dotProduct(ByteSequence a, ByteSequence b); + + /** Returns the sum of squared differences of two signed int8 byte vectors. */ + float squareDistance(ByteSequence a, ByteSequence b); + + /** Returns the cosine similarity of two signed int8 byte vectors. */ + float cosine(ByteSequence a, ByteSequence b); + int hammingDistance(long[] v1, long[] v2); void calculatePartialSums(VectorFloat codebook, int codebookIndex, int size, int clusterCount, VectorFloat query, int offset, VectorSimilarityFunction vsf, VectorFloat partialSums); diff --git a/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java b/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java index 22e0d2c60..f95d4bf6b 100644 --- a/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java +++ b/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java @@ -984,6 +984,37 @@ float assembleAndSumPQ_512( * @param b The second array * @return The Hamming distance */ + @Override + public float dotProduct(ByteSequence a, ByteSequence b) { + float sum = 0; + for (int i = 0; i < a.length(); i++) { + sum += (int) a.get(i) * (int) b.get(i); + } + return sum; + } + + @Override + public float squareDistance(ByteSequence a, ByteSequence b) { + float sum = 0; + for (int i = 0; i < a.length(); i++) { + float diff = a.get(i) - b.get(i); + sum += diff * diff; + } + return sum; + } + + @Override + public float cosine(ByteSequence a, ByteSequence b) { + float dot = 0, normA = 0, normB = 0; + for (int i = 0; i < a.length(); i++) { + float ai = a.get(i), bi = b.get(i); + dot += ai * bi; + normA += ai * ai; + normB += bi * bi; + } + return (float) (dot / Math.sqrt(normA * normB)); + } + @Override public int hammingDistance(long[] a, long[] b) { var sum = LongVector.zero(LongVector.SPECIES_PREFERRED); From 7934fb86e1864fe4d1b8aec0dbc5c839281c7ead Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Thu, 16 Jul 2026 10:19:03 +0000 Subject: [PATCH 3/9] Add ByteVectorSimilarityFunction enum New enum in jvector-base/.../vector/ parallel to VectorSimilarityFunction but operating on ByteSequence, delegating to the VectorUtil byte methods from Sub-Task 2. Three variants with return values normalised to [0,1] matching VectorSimilarityFunction conventions (higher = more similar): EUCLIDEAN: 1 / (1 + squaredL2 / (n * 255^2)) Normalises by the maximum possible squared distance between two signed int8 vectors (255^2 per dimension) so the result stays in (0,1] regardless of dimension. DOT_PRODUCT: (1 + dot / (n * 127^2)) / 2 Normalises by the maximum possible dot product magnitude (127^2 per dimension) before applying the (1+x)/2 mapping so the result stays in [0,1] regardless of dimension or whether vectors are unit-norm. For already unit-norm int8 vectors (e.g. Cohere, OpenAI reduced-precision) prefer COSINE. COSINE: (1 + cosine(v1, v2)) / 2 Cosine is inherently bounded to [-1,1] so no extra normalisation is needed. --- .../vector/ByteVectorSimilarityFunction.java | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/vector/ByteVectorSimilarityFunction.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ByteVectorSimilarityFunction.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ByteVectorSimilarityFunction.java new file mode 100644 index 000000000..2390343ea --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ByteVectorSimilarityFunction.java @@ -0,0 +1,76 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.vector; + +import io.github.jbellis.jvector.vector.types.ByteSequence; + +/** + * Vector similarity function for signed int8 (byte) vectors; parallel to + * {@link VectorSimilarityFunction} but operating on {@link ByteSequence}. + *

+ * Bytes are treated as signed int8 values (Java's {@code byte} is already signed, range −128..127). + * Return-value conventions match {@link VectorSimilarityFunction}: higher is more similar. + */ +public enum ByteVectorSimilarityFunction { + + /** + * Euclidean similarity normalised to {@code (0, 1]}. + * Raw squared L2 is divided by {@code n * 255^2} (the maximum possible squared distance + * between two signed int8 vectors) before the {@code 1 / (1 + x)} mapping, so the result + * is always in (0, 1] regardless of dimension. + */ + EUCLIDEAN { + @Override + public float compare(ByteSequence v1, ByteSequence v2) { + float maxSquaredDist = v1.length() * (255.0f * 255.0f); + return 1.0f / (1.0f + VectorUtil.squareL2Distance(v1, v2) / maxSquaredDist); + } + }, + + /** + * Dot product normalised to {@code [0, 1]}. + * Raw int8 dot product is divided by {@code n * 127^2} (the maximum possible magnitude) + * before applying the {@code (1 + x) / 2} mapping, so the result is always in [0, 1] + * regardless of dimension or whether the vectors are unit-norm. + * For already unit-norm int8 vectors (e.g. Cohere, OpenAI reduced-precision) prefer {@link #COSINE}. + */ + DOT_PRODUCT { + @Override + public float compare(ByteSequence v1, ByteSequence v2) { + float maxMagnitude = v1.length() * (127.0f * 127.0f); + return (1.0f + VectorUtil.dotProduct(v1, v2) / maxMagnitude) / 2.0f; + } + }, + + /** Cosine similarity: {@code (1 + cosine(v1, v2)) / 2} */ + COSINE { + @Override + public float compare(ByteSequence v1, ByteSequence v2) { + return (1.0f + VectorUtil.cosine(v1, v2)) / 2.0f; + } + }; + + /** + * Calculates a similarity score between the two int8 vectors. + * Higher values correspond to closer vectors. + * + * @param v1 a byte vector + * @param v2 another byte vector, of the same dimension + * @return the similarity score + */ + public abstract float compare(ByteSequence v1, ByteSequence v2); +} From 56bc5764e40bece35abb6aebb4f7d0fcf9231e69 Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Thu, 16 Jul 2026 10:19:20 +0000 Subject: [PATCH 4/9] Add BuildScoreProvider.byteVectorScoreProvider factory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire RandomAccessByteVectorValues + ByteVectorSimilarityFunction into the existing BuildScoreProvider abstraction so GraphIndexBuilder can build a graph over byte vectors without any change to the builder core. New static factory: byteVectorScoreProvider(RandomAccessByteVectorValues, ByteVectorSimilarityFunction) The returned BuildScoreProvider: - isExact() -> true; all scoring stays byte×byte with no float round-trip. - approximateCentroid(): sums byte elements cast to float then divides by ravv.size() — acceptable one-time cost per the plan. - searchProviderFor(VectorFloat): throws UnsupportedOperationException with a clear message; float queries not supported on the byte-only build path. - searchProviderFor(int node1): captures ravv.getVector(node1) as 'v', builds an ExactScoreFunction lambda node2 -> bvsf.compare(v, ravv.getVector(node2)), returns a DefaultSearchScoreProvider wrapping it. - diversityProviderFor(int node1): delegates to searchProviderFor(node1). - diversityScoreFunctionFor(int node1): same lambda pattern, returned directly as a ScoreFunction.ExactScoreFunction. Uses the same threadLocalSupplier() pattern as randomAccessScoreProvider() so concurrent graph builds are thread-safe: two independent supplier handles (vectors + vectorsCopy) are created so diversity comparisons don't collide. --- .../graph/similarity/BuildScoreProvider.java | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java index 1049069de..8bec4eb56 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java @@ -16,8 +16,10 @@ package io.github.jbellis.jvector.graph.similarity; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.RemappedRandomAccessVectorValues; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.quantization.BQVectors; import io.github.jbellis.jvector.quantization.PQVectors; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; @@ -211,6 +213,63 @@ public VectorFloat approximateCentroid() { }; } + /** + * Returns a BSP that performs exact score comparisons using the given + * {@link RandomAccessByteVectorValues} and {@link ByteVectorSimilarityFunction}. + * All scoring is byte×byte with no float32 round-trip. + */ + static BuildScoreProvider byteVectorScoreProvider(RandomAccessByteVectorValues ravv, ByteVectorSimilarityFunction bvsf) { + var vectors = ravv.threadLocalSupplier(); + var vectorsCopy = ravv.threadLocalSupplier(); + + return new BuildScoreProvider() { + @Override + public boolean isExact() { + return true; + } + + @Override + public VectorFloat approximateCentroid() { + var vv = vectors.get(); + var centroid = vts.createFloatVector(vv.dimension()); + for (int i = 0; i < vv.size(); i++) { + var v = vv.getVector(i); + for (int d = 0; d < vv.dimension(); d++) { + centroid.set(d, centroid.get(d) + v.get(d)); + } + } + VectorUtil.scale(centroid, 1.0f / vv.size()); + return centroid; + } + + @Override + public SearchScoreProvider searchProviderFor(VectorFloat vector) { + throw new UnsupportedOperationException( + "byteVectorScoreProvider does not support float query vectors; use searchProviderFor(int node)"); + } + + @Override + public SearchScoreProvider searchProviderFor(int node1) { + var v = vectors.get().getVector(node1); + var vc = vectorsCopy.get(); + var sf = (ScoreFunction.ExactScoreFunction) node2 -> bvsf.compare(v, vc.getVector(node2)); + return new DefaultSearchScoreProvider(sf); + } + + @Override + public SearchScoreProvider diversityProviderFor(int node1) { + return searchProviderFor(node1); + } + + @Override + public ScoreFunction diversityScoreFunctionFor(int node1) { + var v = vectors.get().getVector(node1); + var vc = vectorsCopy.get(); + return (ScoreFunction.ExactScoreFunction) node2 -> bvsf.compare(v, vc.getVector(node2)); + } + }; + } + static BuildScoreProvider bqBuildScoreProvider(BQVectors bqv) { return new BuildScoreProvider() { @Override From d5fb1effddf8103287f62b70e50204e38913c301 Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Thu, 16 Jul 2026 10:19:37 +0000 Subject: [PATCH 5/9] GraphIndexBuilder byte-vector constructor, build(), and addGraphNode() overloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit convenience constructor: GraphIndexBuilder(RandomAccessByteVectorValues vectorValues, ByteVectorSimilarityFunction similarityFunction, int M, int beamWidth, float neighborOverflow, float alpha, boolean addHierarchy) Delegates to the BuildScoreProvider constructor via byteVectorScoreProvider(vectorValues, similarityFunction), then stores byteVectorValues and byteVectorSimilarityFunction as nullable instance fields for use by addGraphNode(int, ByteSequence). build(RandomAccessByteVectorValues) overload: Parallel build loop that calls scoreProvider.searchProviderFor(node) directly for each node ordinal, keeping all scoring byte×byte throughout construction. Does not call getVector() and never touches VectorFloat. addGraphNode(int, ByteSequence): Incremental ingest path for byte vectors. Builds an ExactScoreFunction lambda node2 -> bvsf.compare(vector, byteVectorValues.getVector(node2)) and delegates to the existing addGraphNode(int, SearchScoreProvider) — no duplication of graph-construction logic. Guards against misuse (called on a float-only builder) with a clear UnsupportedOperationException pointing to the correct constructor. Two new nullable fields added to GraphIndexBuilder: RandomAccessByteVectorValues byteVectorValues ByteVectorSimilarityFunction byteVectorSimilarityFunction Both are null when the builder is constructed via a float-vector constructor and non-null only when the byte-vector convenience constructor is used. --- .../jvector/graph/GraphIndexBuilder.java | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java index 4139a14b6..188b467dd 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java @@ -24,9 +24,12 @@ import io.github.jbellis.jvector.graph.diversity.VamanaDiversityProvider; import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; import io.github.jbellis.jvector.graph.similarity.ScoreFunction; +import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider; import io.github.jbellis.jvector.util.*; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -74,6 +77,10 @@ public class GraphIndexBuilder implements Closeable, Accountable { private final BuildScoreProvider scoreProvider; + // set only when built from a byte-vector constructor; used by addGraphNode(int, ByteSequence) + private RandomAccessByteVectorValues byteVectorValues; + private ByteVectorSimilarityFunction byteVectorSimilarityFunction; + private final ForkJoinPool simdExecutor; private final ForkJoinPool parallelExecutor; @@ -97,6 +104,39 @@ public class GraphIndexBuilder implements Closeable, Accountable { * an HNSW graph will be created, which is usually not what you want. * @param addHierarchy whether we want to add an HNSW-style hierarchy on top of the Vamana index. */ + /** + * Convenience constructor for building a byte-vector (int8) graph. + * See {@link #GraphIndexBuilder(RandomAccessVectorValues, VectorSimilarityFunction, int, int, float, float, boolean)} + * for the float equivalent. + * + * @param vectorValues the int8 vectors whose relations are represented by the graph + * @param similarityFunction the similarity metric to use during construction + * @param M the maximum number of connections a node can have + * @param beamWidth the size of the beam search to use when finding nearest neighbors + * @param neighborOverflow the ratio of extra neighbors to allow temporarily when inserting a node + * @param alpha how aggressive pruning diverse neighbors should be + * @param addHierarchy whether to add an HNSW-style hierarchy on top of the Vamana index + */ + public GraphIndexBuilder(RandomAccessByteVectorValues vectorValues, + ByteVectorSimilarityFunction similarityFunction, + int M, + int beamWidth, + float neighborOverflow, + float alpha, + boolean addHierarchy) + { + this(BuildScoreProvider.byteVectorScoreProvider(vectorValues, similarityFunction), + vectorValues.dimension(), + M, + beamWidth, + neighborOverflow, + alpha, + addHierarchy, + true); + this.byteVectorValues = vectorValues; + this.byteVectorSimilarityFunction = similarityFunction; + } + public GraphIndexBuilder(RandomAccessVectorValues vectorValues, VectorSimilarityFunction similarityFunction, int M, @@ -446,6 +486,25 @@ public ImmutableGraphIndex build(RandomAccessVectorValues ravv) { cleanup(); return graph; } + + /** + * Builds the graph from a {@link RandomAccessByteVectorValues}. + * Each node is scored via the {@link BuildScoreProvider} supplied at construction time, + * so all comparisons remain byte×byte with no float round-trip. + */ + public ImmutableGraphIndex build(RandomAccessByteVectorValues ravv) { + int size = ravv.size(); + + simdExecutor.submit(() -> { + IntStream.range(0, size).parallel().forEach(node -> { + var ssp = scoreProvider.searchProviderFor(node); + addGraphNode(node, ssp); + }); + }).join(); + + cleanup(); + return graph; + } /** * Validates that the current entry node has been completely added. */ @@ -590,6 +649,26 @@ public long addGraphNode(int node, VectorFloat vector) { return addGraphNode(node, ssp); } + /** + * Inserts a node with the given int8 byte vector into the graph. + * + * @param node the node ID to add + * @param vector the byte vector to add + * @return an estimate of the number of extra bytes used by the graph after adding the given node + * @throws UnsupportedOperationException if this builder was not constructed with a byte-vector score provider + */ + public long addGraphNode(int node, ByteSequence vector) { + if (byteVectorValues == null) { + throw new UnsupportedOperationException( + "addGraphNode(int, ByteSequence) requires a byte-vector GraphIndexBuilder; " + + "use the GraphIndexBuilder(RandomAccessByteVectorValues, ...) constructor"); + } + var bvsf = byteVectorSimilarityFunction; + var ravv = byteVectorValues; + var sf = (ScoreFunction.ExactScoreFunction) node2 -> bvsf.compare(vector, ravv.getVector(node2)); + return addGraphNode(node, new DefaultSearchScoreProvider(sf)); + } + /** * Inserts a node with the given vector value to the graph. * From d3b0e87b5c62934cc6a4fadb13798917104eceeb Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Thu, 16 Jul 2026 10:19:57 +0000 Subject: [PATCH 6/9] test: Fix ambiguous GraphIndexBuilder constructor call in TestVectorGraph Adding GraphIndexBuilder(RandomAccessByteVectorValues, ByteVectorSimilarityFunction, ...) made the existing null,null test call ambiguous since null satisfies both the float and byte-vector overloads. Cast the arguments to (RandomAccessVectorValues) and (VectorSimilarityFunction) to pin the call to the float constructor and restore unambiguous compilation. --- .../java/io/github/jbellis/jvector/graph/TestVectorGraph.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java index 52bdc872a..b784cd7d4 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestVectorGraph.java @@ -433,7 +433,7 @@ public void testGraphIndexBuilderInvalid() { public void testGraphIndexBuilderInvalid(boolean addHierarchy) { assertThrows(NullPointerException.class, - () -> new GraphIndexBuilder(null, null, 0, 0, 1.0f, 1.0f, addHierarchy)); + () -> new GraphIndexBuilder((RandomAccessVectorValues) null, (VectorSimilarityFunction) null, 0, 0, 1.0f, 1.0f, addHierarchy)); // M must be > 0 assertThrows(IllegalArgumentException.class, () -> { From 5a22c1cba4f5334b451fee77d6155a0bd582f83f Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Mon, 20 Jul 2026 07:19:21 +0000 Subject: [PATCH 7/9] Vectorize ByteSequence similarity metrics with SIMD Replace scalar loop implementations of dotProduct, squareDistance, and cosine for ByteSequence with Panama Vector API implementations in PanamaVectorUtilSupport, and native AVX-512/AVX2 kernels wired through NativeVectorUtilSupport -> NativeSimdOps JNI bindings. Java (Panama) path: - Widen signed bytes to int32 via B2I conversion, accumulate products in IntVector lanes, then reduce. - Dispatch on PREFERRED_BIT_SIZE: 512-bit: load 16 bytes (SPECIES_128) -> IntVector.SPECIES_512 256-bit: load 8 bytes (SPECIES_64) -> IntVector.SPECIES_256 128-bit: scalar fallback - cosine variants accumulate dot/norm products in long after reduction to avoid int32 overflow on large vectors. Native path (C++): - New kernels in jvector_simd_kernels.cpp and jvector_avx3_dl_kernels.cpp: dot_product_i8, euclidean_i8, cosine_i8 using Highway SIMD (AVX-512 / AVX2 dispatch). - Registered in jvector_simd_kernel_list.h and exported via jvector_simd.cpp. - Microbenchmarks added in bench_similarity_i8.cpp. - C++ unit tests added in test_similarity_i8.cpp using a prime-length (107-element) vector to exercise tail handling. Java tests: - TestVectorizationProvider.testSimilarityMetricsByte cross-checks SIMD results against scalar DefaultVectorUtilSupport baseline. --- .../vector/NativeVectorUtilSupport.java | 24 ++ .../jvector/vector/cnative/NativeSimdOps.java | 186 +++++++++ .../native/benchmarks/bench_similarity_i8.cpp | 117 ++++++ jvector-native/src/main/native/meson.build | 6 +- .../native/src/jvector_avx3_dl_kernels.cpp | 353 +++++++++++++++++- .../src/main/native/src/jvector_simd.cpp | 9 +- .../native/src/jvector_simd_kernel_list.h | 6 +- .../main/native/src/jvector_simd_kernels.cpp | 169 +++++++++ .../src/main/native/tests/test_helpers.cpp | 22 ++ .../src/main/native/tests/test_helpers.h | 6 +- .../main/native/tests/test_similarity_i8.cpp | 229 ++++++++++++ .../vector/TestVectorizationProvider.java | 35 +- .../vector/PanamaVectorUtilSupport.java | 269 +++++++++++-- 13 files changed, 1397 insertions(+), 34 deletions(-) create mode 100644 jvector-native/src/main/native/benchmarks/bench_similarity_i8.cpp create mode 100644 jvector-native/src/main/native/tests/test_similarity_i8.cpp diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java index 4b627a244..f20ba7805 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java @@ -61,6 +61,30 @@ public String getMaxIsaEnv() { return ptr.reinterpret(Long.MAX_VALUE).getString(0); } + @Override + public float dotProduct(ByteSequence a, ByteSequence b) { + return NativeSimdOps.dot_product_i8( + ((MemorySegmentByteSequence) a).get(), (long) a.offset(), + ((MemorySegmentByteSequence) b).get(), (long) b.offset(), + (long) a.length()); + } + + @Override + public float squareDistance(ByteSequence a, ByteSequence b) { + return NativeSimdOps.euclidean_i8( + ((MemorySegmentByteSequence) a).get(), (long) a.offset(), + ((MemorySegmentByteSequence) b).get(), (long) b.offset(), + (long) a.length()); + } + + @Override + public float cosine(ByteSequence a, ByteSequence b) { + return NativeSimdOps.cosine_i8( + ((MemorySegmentByteSequence) a).get(), (long) a.offset(), + ((MemorySegmentByteSequence) b).get(), (long) b.offset(), + (long) a.length()); + } + @Override protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat vector, int offset) { return FloatVector.fromMemorySegment(SPEC, ((MemorySegmentVectorFloat) vector).get(), vector.offset(offset), ByteOrder.LITTLE_ENDIAN); diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java index d822468a6..19fe50a72 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java @@ -2506,6 +2506,192 @@ public static void nvq_shuffle_query_in_place_8bit(MemorySegment vector, long le } } + private static class dot_product_i8 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("dot_product_i8"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static FunctionDescriptor dot_product_i8$descriptor() { + return dot_product_i8.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MethodHandle dot_product_i8$handle() { + return dot_product_i8.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MemorySegment dot_product_i8$address() { + return dot_product_i8.ADDR; + } + + /** + * {@snippet lang=c : + * float dot_product_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static float dot_product_i8(MemorySegment a, long aoffset, MemorySegment b, long boffset, long length) { + var mh$ = dot_product_i8.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("dot_product_i8", a, aoffset, b, boffset, length); + } + return (float)mh$.invokeExact(a, aoffset, b, boffset, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + + private static class euclidean_i8 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("euclidean_i8"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static FunctionDescriptor euclidean_i8$descriptor() { + return euclidean_i8.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MethodHandle euclidean_i8$handle() { + return euclidean_i8.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MemorySegment euclidean_i8$address() { + return euclidean_i8.ADDR; + } + + /** + * {@snippet lang=c : + * float euclidean_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static float euclidean_i8(MemorySegment a, long aoffset, MemorySegment b, long boffset, long length) { + var mh$ = euclidean_i8.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("euclidean_i8", a, aoffset, b, boffset, length); + } + return (float)mh$.invokeExact(a, aoffset, b, boffset, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + + private static class cosine_i8 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("cosine_i8"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static FunctionDescriptor cosine_i8$descriptor() { + return cosine_i8.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MethodHandle cosine_i8$handle() { + return cosine_i8.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static MemorySegment cosine_i8$address() { + return cosine_i8.ADDR; + } + + /** + * {@snippet lang=c : + * float cosine_i8(const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length) + * } + */ + public static float cosine_i8(MemorySegment a, long aoffset, MemorySegment b, long boffset, long length) { + var mh$ = cosine_i8.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("cosine_i8", a, aoffset, b, boffset, length); + } + return (float)mh$.invokeExact(a, aoffset, b, boffset, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + private static class jvector_simd_get_active_isa { public static final FunctionDescriptor DESC = FunctionDescriptor.of( NativeSimdOps.C_POINTER ); diff --git a/jvector-native/src/main/native/benchmarks/bench_similarity_i8.cpp b/jvector-native/src/main/native/benchmarks/bench_similarity_i8.cpp new file mode 100644 index 000000000..fd9af793a --- /dev/null +++ b/jvector-native/src/main/native/benchmarks/bench_similarity_i8.cpp @@ -0,0 +1,117 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Google Benchmark micro-benchmarks for the int8 vector similarity kernels: +// dot_product_i8, euclidean_i8, cosine_i8 +// +// Parameterised over the realistic embedding dimensions used in production: +// 128, 256, 512, 1024, 1536, 3072 +// +// Build (requires google-benchmark installed or available via pkg-config): +// meson setup build && ninja -C build bench_simd_kernels +// +// Run: +// ./build/bench_simd_kernels [--benchmark_filter=] + +#include +#include +#include + +#include "jvector_simd.h" + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// Deterministic, non-zero int8 vector: values cycle through a signed range to +// avoid degenerate all-zero inputs while staying within [-128, 127]. +static std::vector make_i8_vec(size_t n, int8_t seed) +{ + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + int val = seed + static_cast(i % 127); + if (i % 3 == 0) val = -val; + // clamp to [-127, 127] to keep vectors non-degenerate for cosine + if (val > 127) val = 127; + if (val < -127) val = -127; + v[i] = static_cast(val); + } + return v; +} + +// Benchmark sizes matching production embedding dimensions. +static const std::vector kBenchSizes = {128, 256, 512, 1024, 1536, 3072}; + +// --------------------------------------------------------------------------- +// dot_product_i8 +// --------------------------------------------------------------------------- + +static void BM_dot_product_i8(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_i8_vec(n, 7); + auto b = make_i8_vec(n, 13); + + for (auto _ : state) { + float result = dot_product_i8(a.data(), 0, b.data(), 0, n); + benchmark::DoNotOptimize(result); + } + + state.SetItemsProcessed(state.iterations() * static_cast(n)); + state.SetBytesProcessed(state.iterations() * static_cast(n) * 2 * sizeof(int8_t)); +} +BENCHMARK(BM_dot_product_i8)->ArgsProduct({kBenchSizes}); + +// --------------------------------------------------------------------------- +// euclidean_i8 +// --------------------------------------------------------------------------- + +static void BM_euclidean_i8(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_i8_vec(n, 7); + auto b = make_i8_vec(n, 13); + + for (auto _ : state) { + float result = euclidean_i8(a.data(), 0, b.data(), 0, n); + benchmark::DoNotOptimize(result); + } + + state.SetItemsProcessed(state.iterations() * static_cast(n)); + state.SetBytesProcessed(state.iterations() * static_cast(n) * 2 * sizeof(int8_t)); +} +BENCHMARK(BM_euclidean_i8)->ArgsProduct({kBenchSizes}); + +// --------------------------------------------------------------------------- +// cosine_i8 +// --------------------------------------------------------------------------- + +static void BM_cosine_i8(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_i8_vec(n, 7); + auto b = make_i8_vec(n, 13); + + for (auto _ : state) { + float result = cosine_i8(a.data(), 0, b.data(), 0, n); + benchmark::DoNotOptimize(result); + } + + state.SetItemsProcessed(state.iterations() * static_cast(n)); + state.SetBytesProcessed(state.iterations() * static_cast(n) * 2 * sizeof(int8_t)); +} +BENCHMARK(BM_cosine_i8)->ArgsProduct({kBenchSizes}); + diff --git a/jvector-native/src/main/native/meson.build b/jvector-native/src/main/native/meson.build index 42fec1ada..d992716f3 100644 --- a/jvector-native/src/main/native/meson.build +++ b/jvector-native/src/main/native/meson.build @@ -131,6 +131,7 @@ if gtest_dep.found() sources : [ 'tests/test_helpers.cpp', 'tests/test_similarity.cpp', + 'tests/test_similarity_i8.cpp', 'tests/test_elementwise.cpp', 'tests/test_cpu_features.cpp', ], @@ -154,7 +155,10 @@ gbench_dep = dependency('benchmark', required: false) if gbench_dep.found() executable( 'bench_simd_kernels', - sources : 'benchmarks/bench_similarity_f32.cpp', + sources : [ + 'benchmarks/bench_similarity_f32.cpp', + 'benchmarks/bench_similarity_i8.cpp', + ], dependencies: [vectorutil_dep, gbench_dep], cpp_args : ['-O3'], ) diff --git a/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp b/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp index ae8ab73bb..11e7ed33d 100644 --- a/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp +++ b/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp @@ -25,13 +25,358 @@ // VNNI, VBMI, VBMI2, IFMA, BITALG, VPOPCNTDQ, GFNI, VAES, VPCLMULQDQ // // Compiled with -march=icelake-server. -// Highway will select HWY_AVX3_DL as the static target. +// +// This file uses raw Intel AVX-512 intrinsics directly — NO Google Highway — +// so we get exactly the instructions we intend with zero abstraction overhead. + +#include +#include +#include +#include // AVX-512 + VNNI intrinsics #include "jvector_simd.h" -#include "hwy/highway.h" -#include "assert_hwy_targets.h" -namespace hn = hwy::HWY_NAMESPACE; +// ============================================================================= +// Register naming convention +// zmm = 512-bit (16 × int32, 32 × int16, 64 × int8) +// ymm = 256-bit (32 × int8, 16 × int16) +// xmm = 128-bit (16 × int8, 8 × int16) +// +// VNNI instructions used +// ───────────────────────────────────────────────────────────────────────────── +// VPDPBUSD zmm_acc, zmm_a, zmm_b +// For each group of 4 adjacent lanes (i×4 .. i×4+3): +// acc[i] += (u8)a[i×4+0] * (i8)b[i×4+0] +// + (u8)a[i×4+1] * (i8)b[i×4+1] +// + (u8)a[i×4+2] * (i8)b[i×4+2] +// + (u8)a[i×4+3] * (i8)b[i×4+3] +// → 16 i32 accumulations, 64 int8 products per zmm register per cycle. +// Latency: 3 cycles. Throughput: 1/cycle (two ports on Ice Lake). +// +// VPDPWSSD zmm_acc, zmm_a, zmm_b +// For each group of 2 adjacent i16 lanes (i×2, i×2+1): +// acc[i] += (i16)a[i×2+0] * (i16)b[i×2+0] +// + (i16)a[i×2+1] * (i16)b[i×2+1] +// → 16 i32 accumulations, 32 int16 products per zmm per cycle. +// Latency: 3 cycles. Throughput: 1/cycle. +// +// Signed i8 × signed i8 using VPDPBUSD +// ───────────────────────────────────────────────────────────────────────────── +// VPDPBUSD requires operand A to be unsigned. For signed inputs we apply the +// standard bias trick: +// (a + 128) is always non-negative, so we use it as the unsigned operand. +// (a+128) * b = a*b + 128*b → a*b = VPDPBUSD(a+128, b) - 128 * sum(b) +// +// The bias (128*sum(b)) is constant per zmm load of b, computed as: +// _mm512_dpwssd_epi32(zero, b, set1_epi16(128)) [reuse VPDPWSSD] +// and subtracted once per iteration from the accumulator. +// +// This adds one VPDPWSSD + one VPADDD per iteration, which is negligible +// compared to the main VPDPBUSD throughput. +// +// Unrolling strategy +// ───────────────────────────────────────────────────────────────────────────── +// With 3-cycle VPDPBUSD latency and 1/cycle throughput (ports 0+5), we need +// at least 4 independent accumulator chains to keep the ports saturated: +// issued cycle 0: port 0 ← acc0 +// issued cycle 1: port 5 ← acc1 +// issued cycle 2: port 0 ← acc2 +// issued cycle 3: port 5 ← acc3 (acc0 writeback done, cycle 3) +// 4× unrolling fully hides the 3-cycle latency. +// ============================================================================= namespace AVX3_DL { +// --------------------------------------------------------------------------- +// Horizontal reduce: sum all 16 int32 lanes of a zmm register. +// _mm512_reduce_add_epi32 emits the optimal fold-down sequence; the compiler +// schedules it across surrounding instructions better than manual shuffles. +// --------------------------------------------------------------------------- +static inline int32_t hsum_epi32(__m512i v) +{ + return _mm512_reduce_add_epi32(v); +} + +// --------------------------------------------------------------------------- +// dot_product_i8 — VPDPBUSD with bias correction for signed i8 × signed i8 +// --------------------------------------------------------------------------- +// +// Algorithm +// acc = VPDPBUSD(acc, a_u8, b_i8) where a_u8 = a + 128 +// bias = VPDPWSSD(bias, b_i8, 128) accumulates 128 * sum(b) +// result = hsum(acc) - hsum(bias) +// +// 4× unrolled (256 bytes/iteration) to saturate both ICX VNNI ports and +// fully hide the 3-cycle VPDPBUSD latency. +float dot_product_i8(const int8_t * __restrict__ a, size_t aoffset, + const int8_t * __restrict__ b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + + const __m512i bias128 = _mm512_set1_epi16(128); + + __m512i acc0 = _mm512_setzero_si512(), acc1 = _mm512_setzero_si512(); + __m512i acc2 = _mm512_setzero_si512(), acc3 = _mm512_setzero_si512(); + __m512i bias0 = _mm512_setzero_si512(), bias1 = _mm512_setzero_si512(); + __m512i bias2 = _mm512_setzero_si512(), bias3 = _mm512_setzero_si512(); + + size_t i = 0; + for (; i + 256 <= length; i += 256) { + __m512i va0 = _mm512_loadu_si512(a + i + 0); + __m512i va1 = _mm512_loadu_si512(a + i + 64); + __m512i va2 = _mm512_loadu_si512(a + i + 128); + __m512i va3 = _mm512_loadu_si512(a + i + 192); + __m512i vb0 = _mm512_loadu_si512(b + i + 0); + __m512i vb1 = _mm512_loadu_si512(b + i + 64); + __m512i vb2 = _mm512_loadu_si512(b + i + 128); + __m512i vb3 = _mm512_loadu_si512(b + i + 192); + + // Flip sign bit: maps signed [-128,127] → unsigned [0,255]. + const __m512i flip = _mm512_set1_epi8(-128); + __m512i au0 = _mm512_add_epi8(va0, flip); + __m512i au1 = _mm512_add_epi8(va1, flip); + __m512i au2 = _mm512_add_epi8(va2, flip); + __m512i au3 = _mm512_add_epi8(va3, flip); + + // VPDPBUSD: acc[i] += (u8)au[4i+k] * (i8)vb[4i+k], k=0..3 + acc0 = _mm512_dpbusd_epi32(acc0, au0, vb0); + acc1 = _mm512_dpbusd_epi32(acc1, au1, vb1); + acc2 = _mm512_dpbusd_epi32(acc2, au2, vb2); + acc3 = _mm512_dpbusd_epi32(acc3, au3, vb3); + + // Bias: promote vb to i16 then compute 128 * sum(vb) using VPDPWSSD. + // Each 64-byte zmm of int8 is split into two 512-bit i16 vectors. + __m512i vb0_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 0))); + __m512i vb0_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + __m512i vb1_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 64))); + __m512i vb1_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 96))); + __m512i vb2_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 128))); + __m512i vb2_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 160))); + __m512i vb3_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 192))); + __m512i vb3_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 224))); + + bias0 = _mm512_dpwssd_epi32(bias0, vb0_lo, bias128); + bias0 = _mm512_dpwssd_epi32(bias0, vb0_hi, bias128); + bias1 = _mm512_dpwssd_epi32(bias1, vb1_lo, bias128); + bias1 = _mm512_dpwssd_epi32(bias1, vb1_hi, bias128); + bias2 = _mm512_dpwssd_epi32(bias2, vb2_lo, bias128); + bias2 = _mm512_dpwssd_epi32(bias2, vb2_hi, bias128); + bias3 = _mm512_dpwssd_epi32(bias3, vb3_lo, bias128); + bias3 = _mm512_dpwssd_epi32(bias3, vb3_hi, bias128); + } + __m512i acc = _mm512_add_epi32(_mm512_add_epi32(acc0, acc1), + _mm512_add_epi32(acc2, acc3)); + __m512i bias = _mm512_add_epi32(_mm512_add_epi32(bias0, bias1), + _mm512_add_epi32(bias2, bias3)); + + // Single-zmm tail (residual 64-byte blocks). + for (; i + 64 <= length; i += 64) { + __m512i va = _mm512_loadu_si512(a + i); + __m512i vb = _mm512_loadu_si512(b + i); + __m512i au = _mm512_add_epi8(va, _mm512_set1_epi8(-128)); + acc = _mm512_dpbusd_epi32(acc, au, vb); + + // Promote vb to i16 for bias calculation + __m512i vb_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i vb_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + bias = _mm512_dpwssd_epi32(bias, vb_lo, bias128); + bias = _mm512_dpwssd_epi32(bias, vb_hi, bias128); + } + + int32_t result = hsum_epi32(acc) - hsum_epi32(bias); + + // Scalar tail. + for (; i < length; i++) + result += (int32_t)a[i] * (int32_t)b[i]; + + return (float)result; +} + +// --------------------------------------------------------------------------- +// euclidean_i8 — VPMOVSXBW sign-extend + VPDPWSSD squared differences +// --------------------------------------------------------------------------- +// +// Each 64-byte zmm block is processed as two 32-byte halves: +// _mm512_cvtepi8_epi16(__m256i) = VPMOVSXBW: sign-extends 32×i8 → 32×i16 +// diff = da - db (i16 subtraction, no overflow since range is [-255,255]) +// acc = VPDPWSSD(acc, diff, diff) +// +// 4× unrolled (256 bytes/iteration). +float euclidean_i8(const int8_t * __restrict__ a, size_t aoffset, + const int8_t * __restrict__ b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + + __m512i acc0 = _mm512_setzero_si512(), acc1 = _mm512_setzero_si512(); + __m512i acc2 = _mm512_setzero_si512(), acc3 = _mm512_setzero_si512(); + + size_t i = 0; + for (; i + 256 <= length; i += 256) { +#define EUCL_BLOCK(off, acc_var) \ + { \ + const int8_t *ap = a + i + (off), *bp = b + i + (off); \ + __m512i da_lo = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(ap))); \ + __m512i db_lo = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(bp))); \ + __m512i da_hi = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(ap + 32))); \ + __m512i db_hi = _mm512_cvtepi8_epi16( \ + _mm256_loadu_si256(reinterpret_cast(bp + 32))); \ + __m512i diff_lo = _mm512_sub_epi16(da_lo, db_lo); \ + __m512i diff_hi = _mm512_sub_epi16(da_hi, db_hi); \ + acc_var = _mm512_dpwssd_epi32(acc_var, diff_lo, diff_lo); \ + acc_var = _mm512_dpwssd_epi32(acc_var, diff_hi, diff_hi); \ + } + EUCL_BLOCK( 0, acc0) + EUCL_BLOCK( 64, acc1) + EUCL_BLOCK(128, acc2) + EUCL_BLOCK(192, acc3) +#undef EUCL_BLOCK + } + __m512i acc = _mm512_add_epi32(_mm512_add_epi32(acc0, acc1), + _mm512_add_epi32(acc2, acc3)); + + // Single 64-byte tail blocks. + for (; i + 64 <= length; i += 64) { + __m512i da_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i db_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i da_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 32))); + __m512i db_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + acc = _mm512_dpwssd_epi32(acc, _mm512_sub_epi16(da_lo, db_lo), _mm512_sub_epi16(da_lo, db_lo)); + acc = _mm512_dpwssd_epi32(acc, _mm512_sub_epi16(da_hi, db_hi), _mm512_sub_epi16(da_hi, db_hi)); + } + + // 32-byte tail (one ymm → one 512-bit i16 vector). + if (i + 32 <= length) { + __m512i da = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i db = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i diff = _mm512_sub_epi16(da, db); + acc = _mm512_dpwssd_epi32(acc, diff, diff); + i += 32; + } + + int32_t result = hsum_epi32(acc); + + // Scalar tail. + for (; i < length; i++) { + int32_t d = (int32_t)a[i] - (int32_t)b[i]; + result += d * d; + } + return (float)result; +} + +// --------------------------------------------------------------------------- +// cosine_i8 — three parallel VPDPBUSD chains with bias correction +// --------------------------------------------------------------------------- +// +// Computes dot(a,b), ||a||², ||b||² in a single pass using VPDPBUSD. +// Bias trick: a_u = a+128 (unsigned), then subtract 128*sum(b) and 128*sum(a). +// +// dot(a,b) = hsum(VPDPBUSD(acc_dot, a_u, b)) - 128*sum(b) +// ||a||² = hsum(VPDPBUSD(acc_normA, a_u, a)) - 128*sum(a) +// ||b||² = hsum(VPDPBUSD(acc_normB, b_u, b)) - 128*sum(b) +// +// normB reuses the same biasAB accumulator as dot (both need 128*sum(b)). +// 2× unrolled (128 bytes/iteration) with 6 VPDPBUSD + 4 VPDPWSSD per iter. +float cosine_i8(const int8_t * __restrict__ a, size_t aoffset, + const int8_t * __restrict__ b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + + const __m512i bias128 = _mm512_set1_epi16(128); + const __m512i flip = _mm512_set1_epi8(-128); + + __m512i dot0 = _mm512_setzero_si512(), dot1 = _mm512_setzero_si512(); + __m512i normA0 = _mm512_setzero_si512(), normA1 = _mm512_setzero_si512(); + __m512i normB0 = _mm512_setzero_si512(), normB1 = _mm512_setzero_si512(); + __m512i biasAB0 = _mm512_setzero_si512(), biasAB1 = _mm512_setzero_si512(); + __m512i biasA0 = _mm512_setzero_si512(), biasA1 = _mm512_setzero_si512(); + + size_t i = 0; + for (; i + 128 <= length; i += 128) { + __m512i va0 = _mm512_loadu_si512(a + i); + __m512i vb0 = _mm512_loadu_si512(b + i); + __m512i va1 = _mm512_loadu_si512(a + i + 64); + __m512i vb1 = _mm512_loadu_si512(b + i + 64); + __m512i au0 = _mm512_add_epi8(va0, flip); + __m512i bu0 = _mm512_add_epi8(vb0, flip); + __m512i au1 = _mm512_add_epi8(va1, flip); + __m512i bu1 = _mm512_add_epi8(vb1, flip); + + dot0 = _mm512_dpbusd_epi32(dot0, au0, vb0); + dot1 = _mm512_dpbusd_epi32(dot1, au1, vb1); + normA0 = _mm512_dpbusd_epi32(normA0, au0, va0); + normA1 = _mm512_dpbusd_epi32(normA1, au1, va1); + normB0 = _mm512_dpbusd_epi32(normB0, bu0, vb0); + normB1 = _mm512_dpbusd_epi32(normB1, bu1, vb1); + + // Promote to i16 for bias calculations + __m512i va0_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i va0_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 32))); + __m512i vb0_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i vb0_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + __m512i va1_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 64))); + __m512i va1_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 96))); + __m512i vb1_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 64))); + __m512i vb1_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 96))); + + biasAB0 = _mm512_dpwssd_epi32(biasAB0, vb0_lo, bias128); + biasAB0 = _mm512_dpwssd_epi32(biasAB0, vb0_hi, bias128); + biasAB1 = _mm512_dpwssd_epi32(biasAB1, vb1_lo, bias128); + biasAB1 = _mm512_dpwssd_epi32(biasAB1, vb1_hi, bias128); + biasA0 = _mm512_dpwssd_epi32(biasA0, va0_lo, bias128); + biasA0 = _mm512_dpwssd_epi32(biasA0, va0_hi, bias128); + biasA1 = _mm512_dpwssd_epi32(biasA1, va1_lo, bias128); + biasA1 = _mm512_dpwssd_epi32(biasA1, va1_hi, bias128); + } + __m512i dot = _mm512_add_epi32(dot0, dot1); + __m512i normA = _mm512_add_epi32(normA0, normA1); + __m512i normB = _mm512_add_epi32(normB0, normB1); + __m512i biasAB = _mm512_add_epi32(biasAB0, biasAB1); + __m512i biasA = _mm512_add_epi32(biasA0, biasA1); + + // Single-zmm tail. + for (; i + 64 <= length; i += 64) { + __m512i va = _mm512_loadu_si512(a + i); + __m512i vb = _mm512_loadu_si512(b + i); + __m512i au = _mm512_add_epi8(va, flip); + __m512i bu = _mm512_add_epi8(vb, flip); + dot = _mm512_dpbusd_epi32(dot, au, vb); + normA = _mm512_dpbusd_epi32(normA, au, va); + normB = _mm512_dpbusd_epi32(normB, bu, vb); + + // Promote to i16 for bias calculations + __m512i va_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i))); + __m512i va_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(a + i + 32))); + __m512i vb_lo = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i))); + __m512i vb_hi = _mm512_cvtepi8_epi16(_mm256_loadu_si256(reinterpret_cast(b + i + 32))); + + biasAB = _mm512_dpwssd_epi32(biasAB, vb_lo, bias128); + biasAB = _mm512_dpwssd_epi32(biasAB, vb_hi, bias128); + biasA = _mm512_dpwssd_epi32(biasA, va_lo, bias128); + biasA = _mm512_dpwssd_epi32(biasA, va_hi, bias128); + } + + // Apply bias corrections before scalar tail. + int64_t dotResult = (int64_t)hsum_epi32(dot) - (int64_t)hsum_epi32(biasAB); + int64_t normAResult = (int64_t)hsum_epi32(normA) - (int64_t)hsum_epi32(biasA); + int64_t normBResult = (int64_t)hsum_epi32(normB) - (int64_t)hsum_epi32(biasAB); + + // Scalar tail. + for (; i < length; i++) { + int32_t ai = a[i], bi = b[i]; + dotResult += (int64_t)ai * bi; + normAResult += (int64_t)ai * ai; + normBResult += (int64_t)bi * bi; + } + + return (float)(dotResult / sqrt((double)normAResult * (double)normBResult)); +} + } // namespace AVX3_DL diff --git a/jvector-native/src/main/native/src/jvector_simd.cpp b/jvector-native/src/main/native/src/jvector_simd.cpp index 8bf8ebef5..adc6ba3ac 100644 --- a/jvector-native/src/main/native/src/jvector_simd.cpp +++ b/jvector-native/src/main/native/src/jvector_simd.cpp @@ -82,11 +82,14 @@ static const KernelVTable AVX3_vtable = { }; #undef KERNEL_ENTRY -// AVX3_DL (Ice Lake) inherits all slots from AVX3 unchanged for now. -// To override a slot: t.kernel_name = AVX3_DL::kernel_name; -// The implementation must exist in jvector_avx3_dl_kernels.cpp. +// AVX3_DL (Ice Lake) inherits all slots from AVX3, then overrides the three +// int8 similarity kernels with VNNI-accelerated versions from +// jvector_avx3_dl_kernels.cpp. static const KernelVTable AVX3_DL_vtable = []() { KernelVTable t = AVX3_vtable; + t.dot_product_i8 = AVX3_DL::dot_product_i8; + t.euclidean_i8 = AVX3_DL::euclidean_i8; + t.cosine_i8 = AVX3_DL::cosine_i8; return t; }(); diff --git a/jvector-native/src/main/native/src/jvector_simd_kernel_list.h b/jvector-native/src/main/native/src/jvector_simd_kernel_list.h index 63e7baa4d..99bd99245 100644 --- a/jvector-native/src/main/native/src/jvector_simd_kernel_list.h +++ b/jvector-native/src/main/native/src/jvector_simd_kernel_list.h @@ -58,7 +58,11 @@ KERNEL_ENTRY(float, nvq_square_l2_distance_8bit, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue), (vector, quantized, length, alpha, x0, minValue, maxValue)) \ KERNEL_ENTRY(float, nvq_dot_product_8bit, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue), (vector, quantized, length, alpha, x0, minValue, maxValue)) \ KERNEL_ENTRY(int64_t, nvq_cosine_8bit_packed, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue, const float *centroid), (vector, quantized, length, alpha, x0, minValue, maxValue, centroid)) \ - KERNEL_ENTRY(void, nvq_shuffle_query_in_place_8bit, (float *vector, size_t length), (vector, length)) + KERNEL_ENTRY(void, nvq_shuffle_query_in_place_8bit, (float *vector, size_t length), (vector, length)) \ + /* Int8 byte-vector similarity (VNNI-accelerated on AVX3_DL+) */ \ + KERNEL_ENTRY(float, dot_product_i8, (const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) \ + KERNEL_ENTRY(float, euclidean_i8, (const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) \ + KERNEL_ENTRY(float, cosine_i8, (const int8_t *a, size_t aoffset, const int8_t *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) /* ── ADD NEW KERNEL_ENTRY LINES ABOVE THIS LINE ── */ // clang-format on diff --git a/jvector-native/src/main/native/src/jvector_simd_kernels.cpp b/jvector-native/src/main/native/src/jvector_simd_kernels.cpp index f4e8c2453..1e13dab55 100644 --- a/jvector-native/src/main/native/src/jvector_simd_kernels.cpp +++ b/jvector-native/src/main/native/src/jvector_simd_kernels.cpp @@ -1640,4 +1640,173 @@ HWY_FLATTEN int64_t nvq_cosine_8bit_packed(const float *HWY_RESTRICT vector, return ((int64_t)bmag_bits << 32) | (int64_t)(uint32_t)sum_bits; } +// ============================================================================= +// Int8 byte-vector similarity kernels +// ============================================================================= +// +// These kernels operate on signed int8 (int8_t) vectors — e.g. the output of +// scalar quantization. The generic path here (compiled for SSE4.2, AVX2, AVX3) +// widens i8→i16 using ReorderWidenMulAccumulate, then accumulates into i32. +// +// On the AVX3_DL (Ice Lake+) tier these implementations are overridden in +// jvector_avx3_dl_kernels.cpp with raw AVX-512 VNNI intrinsics — processing +// 64 bytes per VPDPBUSD clock in a single instruction. +// ============================================================================= + +// Horizontal dot product of two signed int8 vectors. +HWY_FLATTEN float dot_product_i8(const int8_t *HWY_RESTRICT a, size_t aoffset, + const int8_t *HWY_RESTRICT b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + const hn::ScalableTag d8; + const hn::RepartitionToWideX2 d32; // int32, lanes = len(d8)/4 + const hn::Repartition d16; // int16 + const size_t lanes8 = hn::Lanes(d8); + + // Four independent accumulators hide the multi-cycle MADD latency. + auto acc0 = hn::Zero(d32), acc1 = hn::Zero(d32); + auto acc2 = hn::Zero(d32), acc3 = hn::Zero(d32); + auto dummy0 = hn::Zero(d32), dummy1 = hn::Zero(d32); + auto dummy2 = hn::Zero(d32), dummy3 = hn::Zero(d32); + size_t i = 0; + for (; i + 4 * lanes8 <= length; i += 4 * lanes8) { + auto va0 = hn::LoadU(d8, a + i); + auto vb0 = hn::LoadU(d8, b + i); + auto va1 = hn::LoadU(d8, a + i + lanes8); + auto vb1 = hn::LoadU(d8, b + i + lanes8); + auto va2 = hn::LoadU(d8, a + i + 2*lanes8); + auto vb2 = hn::LoadU(d8, b + i + 2*lanes8); + auto va3 = hn::LoadU(d8, a + i + 3*lanes8); + auto vb3 = hn::LoadU(d8, b + i + 3*lanes8); + + // Promote to i16 and accumulate using ReorderWidenMulAccumulate (2 i16s -> 1 i32) + acc0 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va0), hn::PromoteLowerTo(d16, vb0), acc0, dummy0); + acc0 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va0), hn::PromoteUpperTo(d16, vb0), acc0, dummy0); + acc1 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va1), hn::PromoteLowerTo(d16, vb1), acc1, dummy1); + acc1 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va1), hn::PromoteUpperTo(d16, vb1), acc1, dummy1); + acc2 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va2), hn::PromoteLowerTo(d16, vb2), acc2, dummy2); + acc2 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va2), hn::PromoteUpperTo(d16, vb2), acc2, dummy2); + acc3 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va3), hn::PromoteLowerTo(d16, vb3), acc3, dummy3); + acc3 = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va3), hn::PromoteUpperTo(d16, vb3), acc3, dummy3); + } + auto acc = hn::Add(hn::Add(acc0, acc1), hn::Add(acc2, acc3)); + auto dummy = hn::Zero(d32); + + for (; i + lanes8 <= length; i += lanes8) { + auto va = hn::LoadU(d8, a + i); + auto vb = hn::LoadU(d8, b + i); + acc = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va), hn::PromoteLowerTo(d16, vb), acc, dummy); + acc = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va), hn::PromoteUpperTo(d16, vb), acc, dummy); + } + int32_t result = hn::ReduceSum(d32, acc); + for (; i < length; i++) result += (int32_t)a[i] * (int32_t)b[i]; + return (float)result; +} + +// Sum of squared differences of two signed int8 vectors. +// Promote i8→i16, subtract in i16, then ReorderWidenMulAccumulate into i32. +HWY_FLATTEN float euclidean_i8(const int8_t *HWY_RESTRICT a, size_t aoffset, + const int8_t *HWY_RESTRICT b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + const hn::ScalableTag d8; + const hn::RepartitionToWideX2 d32; // int32, lanes = len(d8)/4 + const size_t lanes8 = hn::Lanes(d8); + + auto acc0 = hn::Zero(d32), acc0h = hn::Zero(d32); + auto acc1 = hn::Zero(d32), acc1h = hn::Zero(d32); + auto acc2 = hn::Zero(d32), acc2h = hn::Zero(d32); + auto acc3 = hn::Zero(d32), acc3h = hn::Zero(d32); + size_t i = 0; + for (; i + 4 * lanes8 <= length; i += 4 * lanes8) { +#define DO_EUCL_BLOCK(off, lo_var, hi_var) \ + { \ + const hn::RepartitionToWide _d16; \ + auto _va8 = hn::LoadU(d8, a + i + (off)); \ + auto _vb8 = hn::LoadU(d8, b + i + (off)); \ + auto _diff_lo = hn::Sub(hn::PromoteLowerTo(_d16, _va8), \ + hn::PromoteLowerTo(_d16, _vb8)); \ + auto _diff_hi = hn::Sub(hn::PromoteUpperTo(_d16, _va8), \ + hn::PromoteUpperTo(_d16, _vb8)); \ + lo_var = hn::ReorderWidenMulAccumulate(d32, _diff_lo, _diff_lo, lo_var, hi_var); \ + lo_var = hn::ReorderWidenMulAccumulate(d32, _diff_hi, _diff_hi, lo_var, hi_var); \ + } + DO_EUCL_BLOCK(0, acc0, acc0h) + DO_EUCL_BLOCK(lanes8, acc1, acc1h) + DO_EUCL_BLOCK(2*lanes8, acc2, acc2h) + DO_EUCL_BLOCK(3*lanes8, acc3, acc3h) +#undef DO_EUCL_BLOCK + } + auto acc = hn::Add(hn::Add(acc0, acc1), hn::Add(acc2, acc3)); + auto acch = hn::Add(hn::Add(acc0h, acc1h), hn::Add(acc2h, acc3h)); + acc = hn::Add(acc, acch); + for (; i + lanes8 <= length; i += lanes8) { + const hn::RepartitionToWide d16; + auto va8 = hn::LoadU(d8, a + i); + auto vb8 = hn::LoadU(d8, b + i); + auto diff_lo = hn::Sub(hn::PromoteLowerTo(d16, va8), hn::PromoteLowerTo(d16, vb8)); + auto diff_hi = hn::Sub(hn::PromoteUpperTo(d16, va8), hn::PromoteUpperTo(d16, vb8)); + auto dummy_hi = hn::Zero(d32); + acc = hn::ReorderWidenMulAccumulate(d32, diff_lo, diff_lo, acc, dummy_hi); + acc = hn::Add(acc, dummy_hi); + dummy_hi = hn::Zero(d32); + acc = hn::ReorderWidenMulAccumulate(d32, diff_hi, diff_hi, acc, dummy_hi); + acc = hn::Add(acc, dummy_hi); + } + int32_t result = hn::ReduceSum(d32, acc); + for (; i < length; i++) { + int32_t d = (int32_t)a[i] - (int32_t)b[i]; + result += d * d; + } + return (float)result; +} + +// Cosine similarity of two signed int8 vectors. +// Computes dot(a,b), dot(a,a), dot(b,b) in parallel over a single pass. +HWY_FLATTEN float cosine_i8(const int8_t *HWY_RESTRICT a, size_t aoffset, + const int8_t *HWY_RESTRICT b, size_t boffset, + size_t length) +{ + a += aoffset; + b += boffset; + const hn::ScalableTag d8; + const hn::RepartitionToWideX2 d32; + const hn::Repartition d16; + const size_t lanes8 = hn::Lanes(d8); + + auto dot = hn::Zero(d32); + auto normA = hn::Zero(d32); + auto normB = hn::Zero(d32); + auto dummy_acc = hn::Zero(d32); + + size_t i = 0; + for (; i + lanes8 <= length; i += lanes8) { + auto va = hn::LoadU(d8, a + i); + auto vb = hn::LoadU(d8, b + i); + + dot = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va), hn::PromoteLowerTo(d16, vb), dot, dummy_acc); + dot = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va), hn::PromoteUpperTo(d16, vb), dot, dummy_acc); + normA = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, va), hn::PromoteLowerTo(d16, va), normA, dummy_acc); + normA = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, va), hn::PromoteUpperTo(d16, va), normA, dummy_acc); + normB = hn::ReorderWidenMulAccumulate(d32, hn::PromoteLowerTo(d16, vb), hn::PromoteLowerTo(d16, vb), normB, dummy_acc); + normB = hn::ReorderWidenMulAccumulate(d32, hn::PromoteUpperTo(d16, vb), hn::PromoteUpperTo(d16, vb), normB, dummy_acc); + } + + int64_t dotResult = (int64_t)hn::ReduceSum(d32, dot); + int64_t normAResult = (int64_t)hn::ReduceSum(d32, normA); + int64_t normBResult = (int64_t)hn::ReduceSum(d32, normB); + + for (; i < length; i++) { + int32_t ai = a[i], bi = b[i]; + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float)(dotResult / sqrt((double)normAResult * (double)normBResult)); +} + } // namespace JV_ISA diff --git a/jvector-native/src/main/native/tests/test_helpers.cpp b/jvector-native/src/main/native/tests/test_helpers.cpp index 957e42947..45b35cc6e 100644 --- a/jvector-native/src/main/native/tests/test_helpers.cpp +++ b/jvector-native/src/main/native/tests/test_helpers.cpp @@ -73,6 +73,10 @@ const std::vector kKernelTestParams = { {100, "large_mixed_tail"}, {128, "large_power_of_2"}, {255, "large_odd_tail_15"}, + // ---- i8 VNNI 256-byte unroll boundaries (dot_product_i8/euclidean_i8) - + {256, "i8_vnni_4x_exact"}, + {263, "i8_vnni_4x_tail_7"}, + {135, "i8_vnni_2zmm_tail_7"}, }; std::vector make_vec(size_t n, float seed) @@ -86,3 +90,21 @@ std::vector make_vec(size_t n, float seed) return v; } +// Produces n int8_t values with a mix of signs and magnitudes. +// The pattern ensures no element is zero (important for cosine tests). +std::vector make_vec_i8(size_t n, int8_t seed) +{ + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + // Scale seed by a small per-element factor to get variety, + // then clamp to [-100, 100] to keep products well within int16 range. + int val = static_cast(seed) + static_cast(i % 13) - 6; + if (i % 3 == 0) val = -val; // mix of signs + if (val == 0) val = 1; // never zero + if (val > 100) val = 100; + if (val < -100) val = -100; + v[i] = static_cast(val); + } + return v; +} + diff --git a/jvector-native/src/main/native/tests/test_helpers.h b/jvector-native/src/main/native/tests/test_helpers.h index a48ea47cf..e1e9cd79c 100644 --- a/jvector-native/src/main/native/tests/test_helpers.h +++ b/jvector-native/src/main/native/tests/test_helpers.h @@ -35,7 +35,11 @@ // so that no element is exactly zero (important for cosine tests). // --------------------------------------------------------------------------- -std::vector make_vec(size_t n, float seed); +std::vector make_vec(size_t n, float seed); + +// make_vec_i8(n, seed) produces n int8_t values with a mix of signs +// suitable for testing the i8 similarity kernels. +std::vector make_vec_i8(size_t n, int8_t seed); // --------------------------------------------------------------------------- // Shared test parameter — vector length + human-readable path description. diff --git a/jvector-native/src/main/native/tests/test_similarity_i8.cpp b/jvector-native/src/main/native/tests/test_similarity_i8.cpp new file mode 100644 index 000000000..51443a4bd --- /dev/null +++ b/jvector-native/src/main/native/tests/test_similarity_i8.cpp @@ -0,0 +1,229 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Tests for int8 vector similarity kernels: dot_product_i8, euclidean_i8, cosine_i8. +// +// The kernels operate on signed int8_t vectors and return a float result: +// dot_product_i8 — (float) sum(a[i] * b[i]) +// euclidean_i8 — (float) sum((a[i] - b[i])^2) (squared L2 distance) +// cosine_i8 — (float) dot(a,b) / sqrt(||a||^2 * ||b||^2) +// +// On AVX3_DL (Ice Lake+) these are overridden with VNNI (VPDPBUSD/VPDPWSSD) +// implementations; on all other tiers the generic Highway path is used. +// +// All tests are parametrised over kKernelTestParams (defined in test_helpers.cpp), +// which covers every ISA-tier loop-boundary for both f32 and i8 kernels, including +// the VNNI-specific 64/128/256-byte unroll boundaries added for the i8 suite. + +#include "test_helpers.h" + +// --------------------------------------------------------------------------- +// Reference scalar implementations +// --------------------------------------------------------------------------- + +static float ref_dot_i8(const std::vector& a, const std::vector& b) +{ + int64_t s = 0; + for (size_t i = 0; i < a.size(); ++i) + s += static_cast(a[i]) * static_cast(b[i]); + return static_cast(s); +} + +static float ref_euclidean_i8(const std::vector& a, const std::vector& b) +{ + int64_t s = 0; + for (size_t i = 0; i < a.size(); ++i) { + int32_t d = static_cast(a[i]) - static_cast(b[i]); + s += d * d; + } + return static_cast(s); +} + +static float ref_cosine_i8(const std::vector& a, const std::vector& b) +{ + int64_t dot = 0, normA = 0, normB = 0; + for (size_t i = 0; i < a.size(); ++i) { + int32_t ai = a[i], bi = b[i]; + dot += static_cast(ai) * bi; + normA += static_cast(ai) * ai; + normB += static_cast(bi) * bi; + } + return static_cast(dot / std::sqrt(static_cast(normA) + * static_cast(normB))); +} + +// --------------------------------------------------------------------------- +// Parametrised test fixture +// --------------------------------------------------------------------------- + +class SimilarityI8Test : public ::testing::TestWithParam {}; + +// --------------------------------------------------------------------------- +// dot_product_i8 — SIMD result must match the scalar reference +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, DotProduct) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + + const float want = ref_dot_i8(a, b); + const float got = dot_product_i8(a.data(), 0, b.data(), 0, n); + + // Integer accumulation with a single int64→float cast — result is exact. + EXPECT_EQ(got, want); +} + +// --------------------------------------------------------------------------- +// dot_product_i8 with non-zero offsets — exercises the aoffset/boffset path +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, DotProductWithOffset) +{ + const size_t n = GetParam().length; + const size_t prefix = 5; // arbitrary prefix that must be ignored + + std::vector a_pad(prefix + n, 0); + std::vector b_pad(prefix + n, 0); + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + std::copy(a.begin(), a.end(), a_pad.begin() + prefix); + std::copy(b.begin(), b.end(), b_pad.begin() + prefix); + + const float want = ref_dot_i8(a, b); + const float got = dot_product_i8(a_pad.data(), prefix, b_pad.data(), prefix, n); + + EXPECT_EQ(got, want); +} + +// --------------------------------------------------------------------------- +// dot_product_i8 — zero vector gives exactly 0.0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, DotProductZeroVector) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + std::vector z(n, 0); + + EXPECT_EQ(dot_product_i8(a.data(), 0, z.data(), 0, n), 0.0f); +} + +// --------------------------------------------------------------------------- +// euclidean_i8 — SIMD result must match the scalar reference +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, Euclidean) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + + const float want = ref_euclidean_i8(a, b); + const float got = euclidean_i8(a.data(), 0, b.data(), 0, n); + + // Integer accumulation with a single int64→float cast — result is exact. + EXPECT_EQ(got, want); +} + +// --------------------------------------------------------------------------- +// euclidean_i8 — identical vectors must give exactly 0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, EuclideanSameVector) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 9); + + const float got = euclidean_i8(a.data(), 0, a.data(), 0, n); + + EXPECT_EQ(got, 0.0f); +} + +// --------------------------------------------------------------------------- +// cosine_i8 — SIMD result must match the scalar reference +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, Cosine) +{ + const size_t n = GetParam().length; + auto a = make_vec_i8(n, 7); + auto b = make_vec_i8(n, 11); + + const float want = ref_cosine_i8(a, b); + const float got = cosine_i8(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, want, 1e-5f); +} + +// --------------------------------------------------------------------------- +// cosine_i8 — parallel vectors (b = k*a, k > 0) should give similarity ≈ 1.0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, CosineParallelVectors) +{ + const size_t n = GetParam().length; + // Use small magnitudes so that 2*val stays within int8 range. + auto a = make_vec_i8(n, 3); + std::vector b(n); + for (size_t i = 0; i < n; ++i) + b[i] = static_cast(std::max(-127, std::min(127, 2 * static_cast(a[i])))); + + const float got = cosine_i8(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, 1.0f, 1e-5f); +} + +// --------------------------------------------------------------------------- +// cosine_i8 — orthogonal vectors should give similarity ≈ 0.0 +// +// Same analytic construction as the f32 test: for even n, +// a = [+1, +1, +1, ...] +// b = [+1, -1, +1, -1, ...] → dot(a,b) = 0. +// Odd n: the odd last element is zeroed out on b (unchanged on a) so the +// dot product remains zero without affecting the norms materially. +// --------------------------------------------------------------------------- + +TEST_P(SimilarityI8Test, CosineOrthogonalVectors) +{ + const size_t n = GetParam().length; + if (n < 2) GTEST_SKIP() << "need at least 2 elements for orthogonality"; + + const size_t even_n = n - (n % 2); + + std::vector a(n, 0), b(n, 0); + for (size_t i = 0; i < even_n; ++i) { + a[i] = 1; + b[i] = (i % 2 == 0) ? 1 : -1; + } + + const float got = cosine_i8(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, 0.0f, 1e-5f); +} + +// --------------------------------------------------------------------------- +// Instantiation — named using the description field +// --------------------------------------------------------------------------- + +INSTANTIATE_TEST_SUITE_P( + AllSizes, + SimilarityI8Test, + ::testing::ValuesIn(kKernelTestParams), + [](const ::testing::TestParamInfo& info) { + return info.param.description; + }); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java index 81a99aafc..68106cf65 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/vector/TestVectorizationProvider.java @@ -19,7 +19,7 @@ import com.carrotsearch.randomizedtesting.RandomizedTest; import io.github.jbellis.jvector.TestUtil; -import io.github.jbellis.jvector.vector.types.FloatArray; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.junit.Assert; @@ -60,6 +60,39 @@ public void testSimilarityMetricsFloat() { Assert.assertEquals(a.getVectorUtilSupport().squareDistance(v1a, v2a), b.getVectorUtilSupport().squareDistance(v1b, v2b), 0.0001f); } + @Test + public void testSimilarityMetricsByte() { + Assume.assumeTrue(hasSimd); + + VectorizationProvider a = new DefaultVectorizationProvider(); + VectorizationProvider b = VectorizationProvider.getInstance(); + + // Use a prime-length vector that is not a multiple of 8 or 16 + int dim = 107; + byte[] rawA = new byte[dim]; + byte[] rawB = new byte[dim]; + getRandom().nextBytes(rawA); + getRandom().nextBytes(rawB); + + ByteSequence bsA_scalar = a.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence bsB_scalar = a.getVectorTypeSupport().createByteSequence(rawB); + ByteSequence bsA_simd = b.getVectorTypeSupport().createByteSequence(rawA); + ByteSequence bsB_simd = b.getVectorTypeSupport().createByteSequence(rawB); + + Assert.assertEquals( + a.getVectorUtilSupport().dotProduct(bsA_scalar, bsB_scalar), + b.getVectorUtilSupport().dotProduct(bsA_simd, bsB_simd), + 0.0001f); + Assert.assertEquals( + a.getVectorUtilSupport().squareDistance(bsA_scalar, bsB_scalar), + b.getVectorUtilSupport().squareDistance(bsA_simd, bsB_simd), + 0.0001f); + Assert.assertEquals( + a.getVectorUtilSupport().cosine(bsA_scalar, bsB_scalar), + b.getVectorUtilSupport().cosine(bsA_simd, bsB_simd), + 0.0001f); + } + @Test public void testAssembleAndSum() { Assume.assumeTrue(hasSimd); diff --git a/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java b/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java index f95d4bf6b..df49f8858 100644 --- a/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java +++ b/jvector-twenty/src/main/java/io/github/jbellis/jvector/vector/PanamaVectorUtilSupport.java @@ -976,43 +976,266 @@ float assembleAndSumPQ_512( return res; } + // ----------------------------------------------------------------------- + // ByteSequence similarity metrics – Panama SIMD implementations + // + // Strategy: widen signed bytes to int32 via B2I (no AND-mask needed for + // signed arithmetic), accumulate products in IntVector lanes, then reduce. + // The byte-vector species is 1/4 the width of the int species: + // 512-bit int (16 lanes) <- SPECIES_128 bytes + // 256-bit int (8 lanes) <- SPECIES_64 bytes + // 128-bit preferred <- scalar (ByteVector.SPECIES_32 does not exist; + // 128-bit SIMD shows no benefit for this workload) + // ----------------------------------------------------------------------- + /** - * Vectorized calculation of Hamming distance for two arrays of long integers. - * Both arrays should have the same length. - * - * @param a The first array - * @param b The second array - * @return The Hamming distance + * Vectorized dot product of two signed int8 byte vectors. */ @Override public float dotProduct(ByteSequence a, ByteSequence b) { - float sum = 0; - for (int i = 0; i < a.length(); i++) { - sum += (int) a.get(i) * (int) b.get(i); + return switch (PREFERRED_BIT_SIZE) { + case 512 -> dotProductBytes512(a, b); + case 256 -> dotProductBytes256(a, b); + default -> dotProductBytes128(a, b); + }; + } + + float dotProductBytes512(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + final int step = ByteVector.SPECIES_128.length(); // 16 + final int limit = ByteVector.SPECIES_128.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_512); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_128, a, aOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_128, b, bOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + acc = acc.add(va.mul(vb)); } - return sum; + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + result += a.get(aOff + i) * b.get(bOff + i); + } + return result; + } + + float dotProductBytes256(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + final int step = ByteVector.SPECIES_64.length(); // 8 + final int limit = ByteVector.SPECIES_64.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_256); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_64, a, aOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_64, b, bOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + acc = acc.add(va.mul(vb)); + } + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + result += a.get(aOff + i) * b.get(bOff + i); + } + return result; + } + + float dotProductBytes128(ByteSequence a, ByteSequence b) { + // ByteVector.SPECIES_32 does not exist; scalar is fastest at 128-bit width + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + int result = 0; + for (int i = 0; i < length; i++) { + result += a.get(aOff + i) * b.get(bOff + i); + } + return result; } + /** + * Vectorized sum of squared differences between two signed int8 byte vectors. + */ @Override public float squareDistance(ByteSequence a, ByteSequence b) { - float sum = 0; - for (int i = 0; i < a.length(); i++) { - float diff = a.get(i) - b.get(i); - sum += diff * diff; + return switch (PREFERRED_BIT_SIZE) { + case 512 -> squareDistanceBytes512(a, b); + case 256 -> squareDistanceBytes256(a, b); + default -> squareDistanceBytes128(a, b); + }; + } + + float squareDistanceBytes512(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + final int step = ByteVector.SPECIES_128.length(); + final int limit = ByteVector.SPECIES_128.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_512); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_128, a, aOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_128, b, bOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector diff = va.sub(vb); + acc = acc.add(diff.mul(diff)); } - return sum; + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + int diff = a.get(aOff + i) - b.get(bOff + i); + result += diff * diff; + } + return result; + } + + float squareDistanceBytes256(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + final int step = ByteVector.SPECIES_64.length(); + final int limit = ByteVector.SPECIES_64.loopBound(length); + IntVector acc = IntVector.zero(IntVector.SPECIES_256); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_64, a, aOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_64, b, bOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector diff = va.sub(vb); + acc = acc.add(diff.mul(diff)); + } + + int result = acc.reduceLanes(VectorOperators.ADD); + for (int i = limit; i < length; i++) { + int diff = a.get(aOff + i) - b.get(bOff + i); + result += diff * diff; + } + return result; + } + + float squareDistanceBytes128(ByteSequence a, ByteSequence b) { + // ByteVector.SPECIES_32 does not exist; scalar is fastest at 128-bit width + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + int result = 0; + for (int i = 0; i < length; i++) { + int diff = a.get(aOff + i) - b.get(bOff + i); + result += diff * diff; + } + return result; } + /** + * Vectorized cosine similarity between two signed int8 byte vectors. + */ @Override public float cosine(ByteSequence a, ByteSequence b) { - float dot = 0, normA = 0, normB = 0; - for (int i = 0; i < a.length(); i++) { - float ai = a.get(i), bi = b.get(i); - dot += ai * bi; - normA += ai * ai; - normB += bi * bi; - } - return (float) (dot / Math.sqrt(normA * normB)); + return switch (PREFERRED_BIT_SIZE) { + case 512 -> cosineBytes512(a, b); + case 256 -> cosineBytes256(a, b); + default -> cosineBytes128(a, b); + }; + } + + float cosineBytes512(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + final int step = ByteVector.SPECIES_128.length(); + final int limit = ByteVector.SPECIES_128.loopBound(length); + IntVector dot = IntVector.zero(IntVector.SPECIES_512); + IntVector normA = IntVector.zero(IntVector.SPECIES_512); + IntVector normB = IntVector.zero(IntVector.SPECIES_512); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_128, a, aOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_128, b, bOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_512, 0) + .reinterpretAsInts(); + dot = dot.add(va.mul(vb)); + normA = normA.add(va.mul(va)); + normB = normB.add(vb.mul(vb)); + } + + long dotResult = dot.reduceLanes(VectorOperators.ADD); + long normAResult = normA.reduceLanes(VectorOperators.ADD); + long normBResult = normB.reduceLanes(VectorOperators.ADD); + + for (int i = limit; i < length; i++) { + int ai = a.get(aOff + i), bi = b.get(bOff + i); + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float) (dotResult / Math.sqrt((double) normAResult * normBResult)); + } + + float cosineBytes256(ByteSequence a, ByteSequence b) { + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + final int step = ByteVector.SPECIES_64.length(); + final int limit = ByteVector.SPECIES_64.loopBound(length); + IntVector dot = IntVector.zero(IntVector.SPECIES_256); + IntVector normA = IntVector.zero(IntVector.SPECIES_256); + IntVector normB = IntVector.zero(IntVector.SPECIES_256); + + for (int i = 0; i < limit; i += step) { + IntVector va = fromByteSequence(ByteVector.SPECIES_64, a, aOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + IntVector vb = fromByteSequence(ByteVector.SPECIES_64, b, bOff + i) + .convertShape(VectorOperators.B2I, IntVector.SPECIES_256, 0) + .reinterpretAsInts(); + dot = dot.add(va.mul(vb)); + normA = normA.add(va.mul(va)); + normB = normB.add(vb.mul(vb)); + } + + long dotResult = dot.reduceLanes(VectorOperators.ADD); + long normAResult = normA.reduceLanes(VectorOperators.ADD); + long normBResult = normB.reduceLanes(VectorOperators.ADD); + + for (int i = limit; i < length; i++) { + int ai = a.get(aOff + i), bi = b.get(bOff + i); + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float) (dotResult / Math.sqrt((double) normAResult * normBResult)); + } + + float cosineBytes128(ByteSequence a, ByteSequence b) { + // ByteVector.SPECIES_32 does not exist; scalar is fastest at 128-bit width + final int length = a.length(); + final int aOff = a.offset(); + final int bOff = b.offset(); + long dotResult = 0, normAResult = 0, normBResult = 0; + for (int i = 0; i < length; i++) { + int ai = a.get(aOff + i), bi = b.get(bOff + i); + dotResult += ai * bi; + normAResult += ai * ai; + normBResult += bi * bi; + } + return (float) (dotResult / Math.sqrt((double) normAResult * normBResult)); } @Override From 8f7893e881912d570416e0fef1e869b002ed1d5f Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Mon, 17 Aug 2026 03:49:25 +0000 Subject: [PATCH 8/9] Add ScalarQuantizer (VectorCompressor) and SQ tutorial MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ScalarQuantizer implements VectorCompressor> with per-dimension float32→int8 quantization. fit() computes per-dimension min/max from a base RAVV; encode()/encodeTo() map vectors linearly into [-128, 127]; encodeAll() returns an SQVectors in parallel; write()/load() serialize the dimMin/dimMax arrays so the quantizer survives an index reload from disk; dequantize() and quantizeAll() remain as helpers for graph construction and reranking. SQVectors implements CompressedVectors: holds the ByteSequence[] sidecar, provides precomputedScoreFunctionFor/diversityFunctionFor via ByteVectorSimilarityFunction, and has its own write()/load() pair for sidecar persistence. SQExample is a new tutorial that loads siftsmall fvecs, fits a ScalarQuantizer, builds an int8 graph in one parallel shot, saves the graph with original float32 InlineVectors for full-fidelity reranking, and searches every query vector. TutorialRunner gains a "sq" entry point. --- .../jvector/graph/disk/feature/FeatureId.java | 3 +- .../jvector/quantization/SQVectors.java | 148 ++++++++++++ .../jvector/quantization/ScalarQuantizer.java | 212 ++++++++++++++++++ .../jvector/example/tutorial/SQExample.java | 151 +++++++++++++ .../example/tutorial/TutorialRunner.java | 3 + 5 files changed, 516 insertions(+), 1 deletion(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/SQVectors.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ScalarQuantizer.java create mode 100644 jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/SQExample.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java index 131c4c8ee..e9bef0e1d 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FeatureId.java @@ -33,7 +33,8 @@ public enum FeatureId { FUSED_PQ(FusedPQ::load), NVQ_VECTORS(NVQ::load), SEPARATED_VECTORS(SeparatedVectors::load), - SEPARATED_NVQ(SeparatedNVQ::load); + SEPARATED_NVQ(SeparatedNVQ::load), + SQ_QUANTIZER(SQFeature::load); private final BiFunction loader; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/SQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/SQVectors.java new file mode 100644 index 000000000..92d60bbef --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/SQVectors.java @@ -0,0 +1,148 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.quantization; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.similarity.ScoreFunction; +import io.github.jbellis.jvector.util.RamUsageEstimator; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorFloat; + +import java.io.IOException; + +/** + * Compressed vector store produced by {@link ScalarQuantizer}. + * Each vector is stored as a signed int8 {@link ByteSequence}. + * Scoring uses {@link ByteVectorSimilarityFunction} for approximate distance comparisons. + */ +public class SQVectors implements CompressedVectors { + private final ScalarQuantizer sq; + private final ByteSequence[] vectors; + + public SQVectors(ScalarQuantizer sq, ByteSequence[] vectors) { + this.sq = sq; + this.vectors = vectors; + } + + // ── CompressedVectors ──────────────────────────────────────────────────── + + @Override + public void write(IndexWriter out, int version) throws IOException { + sq.write(out, version); + out.writeInt(vectors.length); + for (ByteSequence v : vectors) { + for (int i = 0; i < v.length(); i++) { + out.writeByte(v.get(i)); + } + } + } + + public static SQVectors load(RandomAccessReader in, long offset) throws IOException { + in.seek(offset); + ScalarQuantizer sq = ScalarQuantizer.load(in); + int count = in.readInt(); + int dim = sq.compressedVectorSize(); + ByteSequence[] vectors = new ByteSequence[count]; + byte[] buf = new byte[dim]; + for (int i = 0; i < count; i++) { + in.readFully(buf); + var bs = io.github.jbellis.jvector.vector.VectorizationProvider + .getInstance().getVectorTypeSupport().createByteSequence(dim); + for (int d = 0; d < dim; d++) bs.set(d, buf[d]); + vectors[i] = bs; + } + return new SQVectors(sq, vectors); + } + + @Override + public int getOriginalSize() { + return sq.compressedVectorSize() * Float.BYTES; + } + + @Override + public int getCompressedSize() { + return sq.compressedVectorSize(); + } + + @Override + public ScalarQuantizer getCompressor() { + return sq; + } + + @Override + public int count() { + return vectors.length; + } + + @Override + public long ramBytesUsed() { + if (vectors.length == 0) return 0; + return (long) vectors.length * RamUsageEstimator.sizeOf(new byte[sq.compressedVectorSize()]); + } + + // ── scoring ────────────────────────────────────────────────────────────── + + /** + * Encodes the query with the same {@link ScalarQuantizer} and returns a byte×byte + * approximate score function over the stored int8 vectors. + */ + @Override + public ScoreFunction.ApproximateScoreFunction precomputedScoreFunctionFor( + VectorFloat q, VectorSimilarityFunction similarityFunction) { + return scoreFunctionFor(q, similarityFunction); + } + + @Override + public ScoreFunction.ApproximateScoreFunction scoreFunctionFor( + VectorFloat q, VectorSimilarityFunction similarityFunction) { + ByteVectorSimilarityFunction bvsf = toByteSimFunc(similarityFunction); + ByteSequence qBytes = sq.encode(q); + return node -> bvsf.compare(qBytes, vectors[node]); + } + + @Override + public ScoreFunction.ApproximateScoreFunction diversityFunctionFor( + int node1, VectorSimilarityFunction similarityFunction) { + ByteVectorSimilarityFunction bvsf = toByteSimFunc(similarityFunction); + ByteSequence v1 = vectors[node1]; + return node2 -> bvsf.compare(v1, vectors[node2]); + } + + /** Returns the raw int8 vector for the given node ordinal. */ + public ByteSequence get(int node) { + return vectors[node]; + } + + // ── helpers ────────────────────────────────────────────────────────────── + + private static ByteVectorSimilarityFunction toByteSimFunc(VectorSimilarityFunction vsf) { + switch (vsf) { + case EUCLIDEAN: return ByteVectorSimilarityFunction.EUCLIDEAN; + case DOT_PRODUCT: return ByteVectorSimilarityFunction.DOT_PRODUCT; + case COSINE: return ByteVectorSimilarityFunction.COSINE; + default: throw new IllegalArgumentException("No ByteVectorSimilarityFunction for " + vsf); + } + } + + @Override + public String toString() { + return "SQVectors{count=" + vectors.length + ", sq=" + sq + '}'; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ScalarQuantizer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ScalarQuantizer.java new file mode 100644 index 000000000..0f43bb167 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ScalarQuantizer.java @@ -0,0 +1,212 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.quantization; + +import io.github.jbellis.jvector.disk.IndexWriter; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ListRandomAccessByteVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.ForkJoinPool; +import java.util.stream.IntStream; + +/** + * Per-dimension scalar quantizer: maps float32 vectors to signed int8 using + * per-dimension min/max derived from a base vector set. + * + *

Implements {@link VectorCompressor} so the fitted parameters can be serialized into + * the index header and reloaded when the index is reopened from disk. + * + *

Usage: + * 

+ *   ScalarQuantizer sq = ScalarQuantizer.fit(baseRavv);
+ *   SQVectors sqv       = (SQVectors) sq.encodeAll(baseRavv);
+ *   ByteSequence<?> qb  = sq.encode(queryVector);
+ *   VectorFloat<?>  rec = sq.dequantize(byteVec);
+ *
+ */ +public class ScalarQuantizer implements VectorCompressor> { + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + private final float[] dimMin; + private final float[] dimMax; + + public ScalarQuantizer(float[] dimMin, float[] dimMax) { + this.dimMin = dimMin; + this.dimMax = dimMax; + } + + // ── factory ────────────────────────────────────────────────────────────── + + /** + * Scans all base vectors and computes per-dimension min and max. + */ + public static ScalarQuantizer fit(RandomAccessVectorValues ravv) { + int dim = ravv.dimension(); + float[] dimMin = new float[dim]; + float[] dimMax = new float[dim]; + for (int d = 0; d < dim; d++) { + dimMin[d] = Float.MAX_VALUE; + dimMax[d] = -Float.MAX_VALUE; + } + for (int i = 0; i < ravv.size(); i++) { + VectorFloat v = ravv.getVector(i); + for (int d = 0; d < dim; d++) { + float x = v.get(d); + if (x < dimMin[d]) dimMin[d] = x; + if (x > dimMax[d]) dimMax[d] = x; + } + } + return new ScalarQuantizer(dimMin, dimMax); + } + + // ── VectorCompressor ───────────────────────────────────────────────────── + + /** + * Encodes a single float32 vector to a signed int8 {@link ByteSequence}. + * Each component is mapped linearly from {@code [dimMin[d], dimMax[d]]} to {@code [-128, 127]}. + */ + @Override + public ByteSequence encode(VectorFloat v) { + int dim = dimMin.length; + ByteSequence out = VTS.createByteSequence(dim); + encodeTo(v, out); + return out; + } + + @Override + public void encodeTo(VectorFloat v, ByteSequence dest) { + int dim = dimMin.length; + for (int d = 0; d < dim; d++) { + float range = dimMax[d] - dimMin[d]; + float scaled = range == 0f ? 0f : (v.get(d) - dimMin[d]) / range * 255f - 128f; + int rounded = Math.round(scaled); + dest.set(d, (byte) Math.max(-128, Math.min(127, rounded))); + } + } + + /** Encodes all vectors and returns an {@link SQVectors} instance. */ + @Override + public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + var ravvCopy = ravv.threadLocalSupplier(); + ByteSequence[] encoded = simdExecutor.submit(() -> + IntStream.range(0, ravv.size()) + .parallel() + .mapToObj(i -> { + VectorFloat v = ravvCopy.get().getVector(i); + return v == null ? VTS.createByteSequence(dimMin.length) : encode(v); + }) + .toArray(ByteSequence[]::new) + ).join(); + return new SQVectors(this, encoded); + } + + @Override + @Deprecated + public CompressedVectors createCompressedVectors(Object[] compressedVectors) { + ByteSequence[] seqs = Arrays.copyOf(compressedVectors, compressedVectors.length, ByteSequence[].class); + return new SQVectors(this, seqs); + } + + /** Serialized size of the compressor parameters (dimMin + dimMax arrays). */ + @Override + public int compressorSize() { + // dimension count + two float arrays of that length + return Integer.BYTES + 2 * Float.BYTES * dimMin.length; + } + + /** Each encoded vector is {@code dim} bytes. */ + @Override + public int compressedVectorSize() { + return dimMin.length; + } + + @Override + public void write(IndexWriter out, int version) throws IOException { + out.writeInt(dimMin.length); + for (float v : dimMin) out.writeFloat(v); + for (float v : dimMax) out.writeFloat(v); + } + + /** Deserializes a {@code ScalarQuantizer} written by {@link #write}. */ + public static ScalarQuantizer load(RandomAccessReader in) throws IOException { + int dim = in.readInt(); + float[] dimMin = new float[dim]; + float[] dimMax = new float[dim]; + for (int d = 0; d < dim; d++) dimMin[d] = in.readFloat(); + for (int d = 0; d < dim; d++) dimMax[d] = in.readFloat(); + return new ScalarQuantizer(dimMin, dimMax); + } + + @Override + public double reconstructionError(VectorFloat vector) { + int dim = dimMin.length; + ByteSequence encoded = encode(vector); + VectorFloat reconstructed = dequantize(encoded); + double sum = 0; + for (int d = 0; d < dim; d++) { + double diff = vector.get(d) - reconstructed.get(d); + sum += diff * diff; + } + return sum / dim; + } + + // ── SQ-specific helpers ────────────────────────────────────────────────── + + /** + * Convenience batch-quantize returning a {@link ListRandomAccessByteVectorValues}. + * Useful for constructing the in-memory byte RAVV passed to {@link io.github.jbellis.jvector.graph.GraphIndexBuilder}. + */ + public ListRandomAccessByteVectorValues quantizeAll(RandomAccessVectorValues ravv) { + List> result = new ArrayList<>(ravv.size()); + for (int i = 0; i < ravv.size(); i++) { + result.add(encode(ravv.getVector(i))); + } + return new ListRandomAccessByteVectorValues(result, ravv.dimension()); + } + + /** + * Reconstructs a float32 vector from a signed int8 {@link ByteSequence}. + * Inverse of {@link #encode}: byte -128 → dimMin, byte 127 → dimMax. + */ + public VectorFloat dequantize(ByteSequence b) { + int dim = dimMin.length; + VectorFloat out = VTS.createFloatVector(dim); + for (int d = 0; d < dim; d++) { + float range = dimMax[d] - dimMin[d]; + out.set(d, ((b.get(d) + 128) / 255f) * range + dimMin[d]); + } + return out; + } + + public float[] getDimMin() { return dimMin; } + public float[] getDimMax() { return dimMax; } + + @Override + public String toString() { + return "SQ(per_dim)"; + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/SQExample.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/SQExample.java new file mode 100644 index 000000000..7dda63840 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/SQExample.java @@ -0,0 +1,151 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.example.tutorial; + +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.example.util.SiftLoader; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.GraphSearcher; +import io.github.jbellis.jvector.graph.ImmutableGraphIndex; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; +import io.github.jbellis.jvector.graph.disk.GraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.GraphIndexWriterTypes; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.graph.disk.feature.SQFeature; +import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; +import io.github.jbellis.jvector.graph.similarity.ScoreFunction; +import io.github.jbellis.jvector.quantization.ScalarQuantizer; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.VectorFloat; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; + +/** + * Scalar Quantization (SQ) index build-and-search tutorial using siftsmall float32 vectors. + * + * Loads float32 base and query vectors from .fvecs files, fits a ScalarQuantizer, + * quantizes to int8, builds the graph in one parallel shot, then saves the graph with + * the ScalarQuantizer embedded in the index header via SQFeature. On reload, the + * quantizer is recovered directly from the index — no sidecar file needed. + * + * Run via TutorialRunner: + * ./mvnw -pl jvector-examples -am -Pjdk22 compile exec:exec@tutorial -Dtutorial=sq + */ +public class SQExample { + + /** Siftsmall vectors are 128-dimensional. */ + private static final int DIM = 128; + + public static void main(String[] args) throws IOException { + // ── 1. Load float32 base and query vectors from siftsmall fvecs files ── + String siftPath = "siftsmall"; + List> baseVectors = SiftLoader.readFvecs(siftPath + "/siftsmall_base.fvecs"); + List> queryVectors = SiftLoader.readFvecs(siftPath + "/siftsmall_query.fvecs"); + System.out.printf("Loaded %d base vectors and %d query vectors, dim=%d%n", + baseVectors.size(), queryVectors.size(), DIM); + + // ── 2. Fit a ScalarQuantizer on the base vectors ────────────────────── + // ScalarQuantizer.fit() scans all base vectors to compute per-dimension + // min/max, then maps each float component linearly into [-128, 127]. + var floatRavv = new ListRandomAccessVectorValues(baseVectors, DIM); + ScalarQuantizer sq = ScalarQuantizer.fit(floatRavv); + System.out.println("Quantizer fitted: " + sq); + + // ── 3. Quantize all base vectors to signed int8 ─────────────────────── + RandomAccessByteVectorValues byteRavv = sq.quantizeAll(floatRavv); + System.out.printf("Quantized %d base vectors to int8%n", byteRavv.size()); + + // ── 4. Build the full graph in one shot ─────────────────────────────── + // builder.build(byteRavv) inserts all nodes in parallel and calls cleanup() + // internally — no manual cleanup() needed. + ImmutableGraphIndex graph; + try (GraphIndexBuilder builder = new GraphIndexBuilder( + byteRavv, ByteVectorSimilarityFunction.DOT_PRODUCT, 16, 100, 1.2f, 1.2f, true)) { + graph = builder.build(byteRavv); + } + System.out.printf("Graph built: %d nodes, max level %d%n", + graph.size(0), graph.getMaxLevel()); + + // ── 5. Save the graph to disk ───────────────────────────────────────── + // SQFeature embeds the ScalarQuantizer (dimMin/dimMax) in the index header so + // it is recovered automatically when the index is loaded from disk. + // InlineVectors stores the original float32 vectors for full-fidelity reranking. + Path graphPath = Files.createTempFile("int8-siftsmall", ".jvector"); + try (GraphIndexWriter writer = GraphIndexWriter + .getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, graph, graphPath) + .with(new SQFeature(sq)) + .with(new InlineVectors(DIM)) + .build()) { + writer.write(Map.of( + FeatureId.INLINE_VECTORS, + nodeId -> new InlineVectors.State(floatRavv.getVector(nodeId)) + )); + } + System.out.printf("Graph written to %s (%.1f KB)%n", + graphPath, Files.size(graphPath) / 1024.0); + + // ── 6. Load the graph from disk ─────────────────────────────────────── + // The ScalarQuantizer is recovered directly from the index header — no + // sidecar file and no access to the original base vectors needed. + ReaderSupplier readerSupplier = ReaderSupplierFactory.open(graphPath); + OnDiskGraphIndex diskGraph = OnDiskGraphIndex.load(readerSupplier); + System.out.printf("Graph loaded: %d nodes, max level %d%n", + diskGraph.size(0), diskGraph.getMaxLevel()); + + ScalarQuantizer loadedSq = ((SQFeature) diskGraph.getFeatures().get(FeatureId.SQ_QUANTIZER)) + .getScalarQuantizer(); + System.out.printf("ScalarQuantizer loaded from index header: %s%n", loadedSq); + + // ── 7. Search with every siftsmall query vector ─────────────────────── + // Each float32 query is encoded on-the-fly with the quantizer recovered from disk. + // The search uses byte×byte scoring for graph traversal, then reranks + // the top candidates using the full-precision InlineVectors from disk. + int topK = 10; + int efSearch = 100; + System.out.printf("%nRunning %d queries (topK=%d, efSearch=%d):%n", + queryVectors.size(), topK, efSearch); + + try (GraphSearcher searcher = new GraphSearcher(diskGraph)) { + for (int q = 0; q < queryVectors.size(); q++) { + var queryBytes = loadedSq.encode(queryVectors.get(q)); + var sf = (ScoreFunction.ExactScoreFunction) + node2 -> ByteVectorSimilarityFunction.DOT_PRODUCT.compare(queryBytes, byteRavv.getVector(node2)); + var ssp = new DefaultSearchScoreProvider(sf); + var result = searcher.search(ssp, topK, efSearch, 0.0f, 0.0f, Bits.ALL); + + // Print top-1 result for each query; iterate result.getNodes() for the full list + var top = result.getNodes()[0]; + System.out.printf(" query %3d → top-1 node %5d score %.4f (visited %d nodes)%n", + q, top.node, top.score, result.getVisitedCount()); + } + } + + // ── 8. Cleanup ──────────────────────────────────────────────────────── + readerSupplier.close(); + Files.deleteIfExists(graphPath); + System.out.println("Done."); + } +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java index c675f90c8..011ed306c 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/TutorialRunner.java @@ -41,6 +41,9 @@ public static void main(String[] args) throws IOException { case "nvq": NvqExample.main(forwardArgs); break; + case "sq": + SQExample.main(forwardArgs); + break; default: throw new IllegalArgumentException("Unknown example" + args[0]); } From 4c789ac120a5909266d16449523197d7bddd8a9f Mon Sep 17 00:00:00 2001 From: Raghuveer Devulapalli Date: Thu, 30 Jul 2026 07:56:56 +0000 Subject: [PATCH 9/9] Add SQ compression type to BenchYAML benchmark pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire type: SQ in YAML into the full Grid benchmark path: CompressorParameters.SQParameters is a sentinel that Grid detects with an instanceof check before the normal float32 compressor path, routing to a dedicated buildInt8InMemory() method. buildInt8InMemory() fits a ScalarQuantizer on the base RAVV, quantizes all base vectors to int8, builds the graph via GraphIndexBuilder(RandomAccessByteVectorValues, ByteVectorSimilarityFunction), embeds the ScalarQuantizer in the index header via SQFeature (so it is recoverable at search time without access to base vectors), and writes original float32 vectors as INLINE_VECTORS for full-fidelity reranking. ConfiguredSystem gains a byteRavv field and an SQ-aware constructor. scoreProviderFor() detects SQ_QUANTIZER in the feature set, recovers the ScalarQuantizer from the OnDiskGraphIndex header, encodes the query on- the-fly, uses byte×byte scoring as the ApproximateScoreFunction for graph traversal, and falls back to the float INLINE_VECTORS reranker for final topK selection. Also adds sift-128-euclidean-int8.yml benchmark config and a corresponding datasets.yml entry. --- .../github/jbellis/jvector/example/Grid.java | 171 +++++++++++++++++- .../example/util/CompressorParameters.java | 23 ++- .../jvector/example/yaml/Compression.java | 2 + jvector-examples/yaml-configs/datasets.yml | 5 +- .../sift-128-euclidean-int8.yml | 27 +++ 5 files changed, 224 insertions(+), 4 deletions(-) create mode 100644 jvector-examples/yaml-configs/index-parameters/sift-128-euclidean-int8.yml diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index 8f45df2a0..c171b8b21 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -35,10 +35,13 @@ import io.github.jbellis.jvector.example.util.CompressorParameters; import io.github.jbellis.jvector.example.util.FilteredForkJoinPool; import io.github.jbellis.jvector.example.util.OnDiskGraphIndexCache; +import io.github.jbellis.jvector.quantization.ScalarQuantizer; import io.github.jbellis.jvector.example.yaml.MetricSelection; import io.github.jbellis.jvector.graph.ImmutableGraphIndex; import io.github.jbellis.jvector.graph.GraphIndexBuilder; import io.github.jbellis.jvector.graph.GraphSearcher; +import io.github.jbellis.jvector.graph.ListRandomAccessByteVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessByteVectorValues; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.*; import io.github.jbellis.jvector.graph.disk.feature.Feature; @@ -46,6 +49,7 @@ import io.github.jbellis.jvector.graph.disk.feature.FusedPQ; import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; import io.github.jbellis.jvector.graph.disk.feature.NVQ; +import io.github.jbellis.jvector.graph.disk.feature.SQFeature; import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; import io.github.jbellis.jvector.graph.similarity.ScoreFunction; @@ -57,6 +61,9 @@ import io.github.jbellis.jvector.quantization.VectorCompressor; import io.github.jbellis.jvector.util.ExplicitThreadLocal; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.ByteSequence; import io.github.jbellis.jvector.vector.types.VectorFloat; import java.io.FileNotFoundException; @@ -242,8 +249,48 @@ static void runOneGraph(OnDiskGraphIndexCache cache, VectorCompressor buildCompressorObj = null; String buildQuantType = null; - if (buildCompressor != null) { - var buildParams = buildCompressor.apply(ds); + // Check for INT8/SQ path before resolving the compressor object + CompressorParameters buildParams = buildCompressor != null ? buildCompressor.apply(ds) : null; + + if (buildParams instanceof CompressorParameters.SQParameters) { + // --- INT8 scalar-quantization build path --- + String buildCompressorString = "SQ(per_dim)"; + ScalarQuantizer sq = ScalarQuantizer.fit(ds.getBaseRavv()); + Int8BuildResult int8Result = + buildInt8InMemory(featureSets, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, ds, sq, workDirectory); + + // Capture post-build metrics + diagnostics.capturePostPhaseSnapshot("Graph Build"); + diagnostics.printDiskStatistics("Graph Index Build"); + System.out.printf("Index build time: %f seconds%n%n", Grid.getIndexBuildTimeSeconds(ds.getName())); + constructionMetrics.indexBuildTimeS = Grid.getIndexBuildTimeSeconds(ds.getName()); + + try { + int8Result.indexes.forEach((features, index) -> { + final Set featureSetForIndex = index instanceof OnDiskGraphIndex + ? ((OnDiskGraphIndex) index).getFeatureSet() : Set.of(); + try (var cs = new ConfiguredSystem(ds, index, null, featureSetForIndex, int8Result.byteRavv)) { + testConfiguration(cs, topKGrid, usePruningGrid, M, efConstruction, neighborOverflow, addHierarchy, refineFinalGraph, + featureSetForIndex, buildCompressorString, artifacts, constructionMetrics, workDirectory); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + for (var index : int8Result.indexes.values()) { + index.close(); + } + } finally { + for (int nn = 0; nn < featureSets.size(); nn++) { + Path p = workDirectory.resolve("graph" + nn); + try { Files.deleteIfExists(p); } catch (IOException e) { + System.err.println("Cleanup Failed: Could not delete " + p.getFileName() + " -> " + e.getMessage()); + } + } + } + return; // done for this grid cell + } + + if (buildParams != null) { buildQuantType = quantTypeOf(buildParams); // "PQ", "BQ", or null buildCompressorObj = getCompressor(buildCompressor, ds, constructionMetrics, Phase.INDEX, buildQuantType); } @@ -608,6 +655,95 @@ private static Map, ImmutableGraphIndex> buildInMemory(List, ImmutableGraphIndex> indexes; + final ListRandomAccessByteVectorValues byteRavv; + Int8BuildResult(Map, ImmutableGraphIndex> indexes, ListRandomAccessByteVectorValues byteRavv) { + this.indexes = indexes; + this.byteRavv = byteRavv; + } + } + + /** + * Builds an INT8 graph index in-memory from a float32 DataSet by applying per-dimension + * scalar quantization. Embeds the ScalarQuantizer in the index header via SQFeature so it + * is available at search time without re-fitting. Writes original float32 vectors as + * INLINE_VECTORS for full-fidelity reranking. + * Feature sets other than INLINE_VECTORS are skipped for INT8 builds. + */ + private static Int8BuildResult buildInt8InMemory(List> featureSets, + int M, + int efConstruction, + float neighborOverflow, + boolean addHierarchy, + boolean refineFinalGraph, + DataSet ds, + ScalarQuantizer sq, + Path testDirectory) + throws IOException + { + var floatVectors = ds.getBaseRavv(); + System.out.format("%s: Scalar-quantizing %d vectors (per_dim)%n", ds.getName(), floatVectors.size()); + long sqStart = System.nanoTime(); + ListRandomAccessByteVectorValues byteRavv = sq.quantizeAll(floatVectors); + System.out.format("%s: Quantization done in %.2fs%n", ds.getName(), (System.nanoTime() - sqStart) / 1e9); + + ByteVectorSimilarityFunction byteSimFunc = toByteSimFunc(ds.getSimilarityFunction()); + + GraphIndexBuilder builder = new GraphIndexBuilder(byteRavv, byteSimFunc, M, efConstruction, neighborOverflow, 1.2f, addHierarchy); + long start = System.nanoTime(); + var onHeapGraph = builder.build(byteRavv); + double buildTimeS = (System.nanoTime() - start) / 1_000_000_000.0; + System.out.format("Build (INT8/SQ) M=%d overflow=%.2f ef=%d in %.2fs%n", M, neighborOverflow, efConstruction, buildTimeS); + for (int i = 0; i <= onHeapGraph.getMaxLevel(); i++) { + System.out.format(" L%d: %d nodes, %.2f avg degree%n", i, onHeapGraph.size(i), onHeapGraph.getAverageDegree(i)); + } + + // Validate up-front: every feature set must contain INLINE_VECTORS. + // NVQ and FUSED_PQ require float32 codebooks and are not supported for INT8 builds. + for (var features : featureSets) { + if (!features.contains(FeatureId.INLINE_VECTORS)) { + throw new IllegalArgumentException( + "INT8/SQ build requires reranking: [FP] in YAML (feature set must contain INLINE_VECTORS). " + + "Got: " + features + ". NVQ and FUSED_PQ are not supported for INT8 builds."); + } + } + + Map, ImmutableGraphIndex> indexes = new HashMap<>(); + int n = 0; + for (var features : featureSets) { + var graphPath = testDirectory.resolve("graph" + n++); + var identityMapper = new OrdinalMapper.IdentityMapper(byteRavv.size() - 1); + var writer = new OnDiskGraphIndexWriter.Builder(onHeapGraph, graphPath) + .withMapper(identityMapper) + .with(new SQFeature(sq)) + .with(new InlineVectors(floatVectors.dimension())) + .build(); + try (writer) { + start = System.nanoTime(); + writer.write(Map.of( + FeatureId.INLINE_VECTORS, + (IntFunction) nodeId -> new InlineVectors.State(floatVectors.getVector(nodeId)) + )); + System.out.format("Wrote %s (INT8) in %.2fs%n", features, (System.nanoTime() - start) / 1_000_000_000.0); + } + indexes.put(features, OnDiskGraphIndex.load(ReaderSupplierFactory.open(graphPath))); + } + indexBuildTimes.put(ds.getName(), buildTimeS); + return new Int8BuildResult(indexes, byteRavv); + } + + /** Maps a float VectorSimilarityFunction to its byte equivalent. */ + private static ByteVectorSimilarityFunction toByteSimFunc(VectorSimilarityFunction vsf) { + switch (vsf) { + case EUCLIDEAN: return ByteVectorSimilarityFunction.EUCLIDEAN; + case DOT_PRODUCT: return ByteVectorSimilarityFunction.DOT_PRODUCT; + case COSINE: return ByteVectorSimilarityFunction.COSINE; + default: throw new IllegalArgumentException("No ByteVectorSimilarityFunction for " + vsf); + } + } + // avoid recomputing the compressor repeatedly (this is a relatively small memory footprint) static final Map> cachedCompressors = new IdentityHashMap<>(); @@ -1101,18 +1237,49 @@ public static class ConfiguredSystem implements AutoCloseable { CompressedVectors cv; Set features; + // Non-null for INT8/SQ builds; null otherwise. sq is recovered from the index header. + final ListRandomAccessByteVectorValues byteRavv; + private final ExplicitThreadLocal searchers = ExplicitThreadLocal.withInitial(() -> { return new GraphSearcher(index); }); + /** Constructor for float32 builds. */ ConfiguredSystem(DataSet ds, ImmutableGraphIndex index, CompressedVectors cv, Set features) { this.ds = ds; this.index = index; this.cv = cv; this.features = features; + this.byteRavv = null; + } + + /** Constructor for INT8/SQ builds. sq is recovered from the SQFeature in the index header. */ + ConfiguredSystem(DataSet ds, ImmutableGraphIndex index, CompressedVectors cv, + Set features, ListRandomAccessByteVectorValues byteRavv) { + this.ds = ds; + this.index = index; + this.cv = cv; + this.features = features; + this.byteRavv = byteRavv; } public SearchScoreProvider scoreProviderFor(VectorFloat queryVector, ImmutableGraphIndex.View view) { + // INT8/SQ path: recover the ScalarQuantizer from the index header (SQFeature), + // encode the query on-the-fly, score byte×byte against the in-memory byteRavv, + // then rerank via INLINE_VECTORS (original float32 vectors) for final accuracy. + if (features.contains(FeatureId.SQ_QUANTIZER)) { + ScalarQuantizer sq = ((SQFeature) ((OnDiskGraphIndex) index).getFeatures().get(FeatureId.SQ_QUANTIZER)) + .getScalarQuantizer(); + ByteSequence qByte = sq.encode(queryVector); + ByteVectorSimilarityFunction byteSimFunc = toByteSimFunc(ds.getSimilarityFunction()); + ScoreFunction.ApproximateScoreFunction asf = + node -> byteSimFunc.compare(qByte, byteRavv.getVector(node)); + var scoringView = (ImmutableGraphIndex.ScoringView) view; + var rr = scoringView.rerankerFor(queryVector, ds.getSimilarityFunction()); + return new DefaultSearchScoreProvider(asf, rr); + } + + // Float32 path (unchanged) var scoringView = (ImmutableGraphIndex.ScoringView) view; ScoreFunction.ApproximateScoreFunction asf; if (features.contains(FeatureId.FUSED_PQ)) { diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/CompressorParameters.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/CompressorParameters.java index 2f4aceaf7..022e86f8e 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/CompressorParameters.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/CompressorParameters.java @@ -21,7 +21,6 @@ import io.github.jbellis.jvector.quantization.NVQuantization; import io.github.jbellis.jvector.quantization.ProductQuantization; import io.github.jbellis.jvector.quantization.VectorCompressor; - public abstract class CompressorParameters { public static final CompressorParameters NONE = new NoCompressionParameters(); @@ -95,6 +94,28 @@ public boolean supportsCaching() { } } + /** + * Sentinel parameters for scalar (INT8) quantization. + * Grid detects this type via instanceof and routes to the INT8 build path; + * computeCompressor() is never called on this class. + */ + public static class SQParameters extends CompressorParameters { + @Override + public VectorCompressor computeCompressor(DataSet ds) { + throw new UnsupportedOperationException("SQ uses the INT8 build path; computeCompressor should not be called"); + } + + @Override + public String idStringFor(DataSet ds) { + return "SQ_per_dim_" + ds.getName(); + } + + @Override + public boolean supportsCaching() { + return false; + } + } + private static class NoCompressionParameters extends CompressorParameters { @Override public VectorCompressor computeCompressor(DataSet ds) { diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java index a8277508a..268a60da6 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/yaml/Compression.java @@ -60,6 +60,8 @@ public Function getCompressorParameters() { }; case "BQ": return ds -> new CompressorParameters.BQParameters(); + case "SQ": + return ds -> new CompressorParameters.SQParameters(); default: throw new IllegalArgumentException("Unsupported compression type: " + type); diff --git a/jvector-examples/yaml-configs/datasets.yml b/jvector-examples/yaml-configs/datasets.yml index 60aa893f5..b9fb72007 100644 --- a/jvector-examples/yaml-configs/datasets.yml +++ b/jvector-examples/yaml-configs/datasets.yml @@ -1,3 +1,6 @@ +int8-benchmarks: + - sift-128-euclidean-int8 + jvector-100k: - cohere-english-v3-100k - ada002-100k @@ -37,4 +40,4 @@ openai-unconfigured: # - cohere-english-v3-1M # - cohere-english-v3-10M # - deep-image-96-angular # large files not yet supported -# - gist-960-euclidean # large files not yet supported \ No newline at end of file +# - gist-960-euclidean # large files not yet supported diff --git a/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean-int8.yml b/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean-int8.yml new file mode 100644 index 000000000..8f2c66c30 --- /dev/null +++ b/jvector-examples/yaml-configs/index-parameters/sift-128-euclidean-int8.yml @@ -0,0 +1,27 @@ +yamlSchemaVersion: 1 +onDiskIndexVersion: 6 + +dataset: sift-128-euclidean + +construction: + outDegree: [32] + efConstruction: [100] + neighborOverflow: [1.2f] + addHierarchy: [Yes] + refineFinalGraph: [Yes] + fusedGraph: [No] + compression: + - type: SQ + # Per-dimension min/max scalar quantization: maps float32 → int8. + # No parameters required; scheme is fixed to per_dim. + reranking: + - FP # INLINE_VECTORS stores dequantized float32 for reranking after INT8 graph traversal + useSavedIndexIfExists: No + +search: + topKOverquery: + 10: [1.0, 2.0, 5.0, 10.0] + 100: [1.0, 2.0] + useSearchPruning: [Yes] + compression: + - type: None # INT8 search uses byte scoring directly via ByteVectorSimilarityFunction