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. * 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; + } +} 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/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 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-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); +} 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-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/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]); } 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 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/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, () -> { 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 22e0d2c60..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,14 +976,268 @@ 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 dot product of two signed int8 byte vectors. + */ + @Override + public float dotProduct(ByteSequence a, ByteSequence b) { + 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)); + } + + 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 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 sum of squared differences between two signed int8 byte vectors. */ + @Override + public float squareDistance(ByteSequence a, ByteSequence b) { + 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)); + } + + 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) { + 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 public int hammingDistance(long[] a, long[] b) { var sum = LongVector.zero(LongVector.SPECIES_PREFERRED);