diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml index c75b6f39b..0b6bb3322 100644 --- a/.github/workflows/unit-tests.yaml +++ b/.github/workflows/unit-tests.yaml @@ -45,12 +45,38 @@ jobs: - name: Set up GCC run: | sudo apt install -y gcc - - name: Install Meson and Ninja + - name: Install Meson, Ninja, and GTest run: | - sudo apt update && sudo apt install -y meson ninja-build + sudo apt update && sudo apt install -y meson ninja-build pkg-config libgtest-dev - uses: actions/checkout@v4 - name: Initialize Git Submodules run: git submodule update --init + + - name: Build test_simd_kernels (native C++) + working-directory: jvector-native/src/main/native + run: | + meson setup build --wipe + ninja -C build test_simd_kernels + + - name: Run test_simd_kernels — no ISA cap (auto-detect) + if: matrix.max_isa == 'avx512f' + working-directory: jvector-native/src/main/native + run: ./build/test_simd_kernels + + - name: Run test_simd_kernels — capped at avx2 + if: matrix.max_isa == 'avx2' + working-directory: jvector-native/src/main/native + env: + JVECTOR_MAX_ISA: avx2 + run: ./build/test_simd_kernels + + - name: Run test_simd_kernels — capped at sse42 + if: matrix.max_isa == 'sse42' + working-directory: jvector-native/src/main/native + env: + JVECTOR_MAX_ISA: sse42 + run: ./build/test_simd_kernels + - name: Set up JDK ${{ matrix.jdk }} uses: actions/setup-java@v3 with: 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/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/ScalarQuantizer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ScalarQuantizer.java new file mode 100644 index 000000000..92ce92e45 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ScalarQuantizer.java @@ -0,0 +1,124 @@ +/* + * 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.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.util.ArrayList; +import java.util.List; + +/** + * Per-dimension scalar quantizer: maps float32 vectors to signed int8 using + * per-dimension min/max derived from a base vector set. + * + *

Usage: + * 

+ *    ScalarQuantizer sq = ScalarQuantizer.fit(baseRavv);
+ *    ListRandomAccessByteVectorValues byteRavv = sq.quantizeAll(baseRavv);
+ *    ByteSequence<?> queryBytes = sq.quantize(queryVector);
+ *    VectorFloat<?> reranked   = sq.dequantize(byteVec);
+ *
+ */ +public class ScalarQuantizer { + private static final VectorTypeSupport VTS = + VectorizationProvider.getInstance().getVectorTypeSupport(); + + private final float[] dimMin; + private final float[] dimMax; + + private ScalarQuantizer(float[] dimMin, float[] dimMax) { + this.dimMin = dimMin; + this.dimMax = dimMax; + } + + /** + * 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); + } + + /** + * Quantizes a single float32 vector to a signed int8 ByteSequence. + * Each component is mapped linearly from [dimMin[d], dimMax[d]] to [-128, 127]. + */ + public ByteSequence quantize(VectorFloat v) { + int dim = dimMin.length; + ByteSequence out = VTS.createByteSequence(dim); + 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); + out.set(d, (byte) Math.max(-128, Math.min(127, rounded))); + } + return out; + } + + /** + * Quantizes all vectors in the given RAVV and returns them as a + * {@link ListRandomAccessByteVectorValues}. + */ + public ListRandomAccessByteVectorValues quantizeAll(RandomAccessVectorValues ravv) { + List> result = new ArrayList<>(ravv.size()); + for (int i = 0; i < ravv.size(); i++) { + result.add(quantize(ravv.getVector(i))); + } + return new ListRandomAccessByteVectorValues(result, ravv.dimension()); + } + + /** + * Reconstructs a float32 vector from a signed int8 ByteSequence using the stored + * per-dimension parameters. Used to populate INLINE_VECTORS for reranking. + * + *

Inverse of {@link #quantize}: byte -128 → dimMin, byte 127 → dimMax. + * Uses {@code (b + 128)} to undo the [-128, 127] shift applied during quantization. + */ + 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; + } + + @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..52c1cb43a 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; @@ -57,6 +60,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 +248,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, sq, int8Result.byteRavv, featureSetForIndex)) { + 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 +654,92 @@ 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, then writing INLINE_VECTORS (dequantized float32) for 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 InlineVectors(floatVectors.dimension())) + .build(); + try (writer) { + start = System.nanoTime(); + writer.write(Map.of( + FeatureId.INLINE_VECTORS, + (IntFunction) nodeId -> new InlineVectors.State(sq.dequantize(byteRavv.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 +1233,51 @@ public static class ConfiguredSystem implements AutoCloseable { CompressedVectors cv; Set features; + // Non-null for INT8/SQ builds; null for float32 builds + final ScalarQuantizer sq; + 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.sq = null; + this.byteRavv = null; + } + + /** Constructor for INT8/SQ builds. cv is always null; sq and byteRavv drive scoring. */ + ConfiguredSystem(DataSet ds, ImmutableGraphIndex index, ScalarQuantizer sq, + ListRandomAccessByteVectorValues byteRavv, Set features) { + this.ds = ds; + this.index = index; + this.cv = null; + this.features = features; + this.sq = sq; + this.byteRavv = byteRavv; } public SearchScoreProvider scoreProviderFor(VectorFloat queryVector, ImmutableGraphIndex.View view) { + // INT8 path: quantize query on-the-fly; score byte×byte as an approximation; + // rerank via INLINE_VECTORS (dequantized float32) for final accuracy. + // We declare it as ApproximateScoreFunction so GraphSearcher treats the byte scores + // as a proxy ranking and always applies the float reranker to the top rerankK candidates. + if (sq != null) { + ByteSequence qByte = sq.quantize(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/Int8IndexBuild.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/Int8IndexBuild.java new file mode 100644 index 000000000..962ca7f3c --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/tutorial/Int8IndexBuild.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.ListRandomAccessByteVectorValues; +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.similarity.DefaultSearchScoreProvider; +import io.github.jbellis.jvector.graph.similarity.ScoreFunction; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.ByteVectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; + +/** + * Int8 HNSW/Vamana index build-and-search tutorial using real siftsmall int8 vectors. + * + * Reads pre-quantised int8 base and query vectors from the siftsmall directory, + * builds an on-disk graph index, then runs every query vector against it. + * + * Run via TutorialRunner: + * ./mvnw -pl jvector-examples -am -Pjdk22 compile exec:exec@tutorial -Dtutorial=int8build + */ +public class Int8IndexBuild { + + /** Siftsmall vectors are 128-dimensional. */ + private static final int DIM = 128; + + public static void main(String[] args) throws IOException { + // ── 1. Vector type support ───────────────────────────────────────────── + VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + System.out.println("VectorizationProvider: " + VectorizationProvider.getInstance().getClass().getSimpleName()); + + // ── 2. Load int8 vectors from siftsmall directory ───────────────────── + // The .bvecs files are raw flat binary: numVectors * DIM bytes, no header. + String siftPath = "siftsmall"; + List> baseVectors = SiftLoader.readBvecs(siftPath + "/siftsmall_base.bvecs", DIM); + List> queryVectors = SiftLoader.readBvecs(siftPath + "/siftsmall_query.bvecs", DIM); + System.out.printf("Loaded %d base vectors and %d query vectors, dim=%d%n", + baseVectors.size(), queryVectors.size(), DIM); + + // ── 3. Wrap base vectors in a RandomAccessByteVectorValues (RAVV) ───── + RandomAccessByteVectorValues ravv = new ListRandomAccessByteVectorValues(baseVectors, DIM); + + // ── 4. Build the graph incrementally (int8, DOT_PRODUCT, M=16, efConstruction=100) ─ + ImmutableGraphIndex graph; + try (GraphIndexBuilder builder = new GraphIndexBuilder( + ravv, ByteVectorSimilarityFunction.DOT_PRODUCT, 16, 100, 1.2f, 1.2f, true)) { + + for (int i = 0; i < baseVectors.size(); i++) { + builder.addGraphNode(i, baseVectors.get(i)); + } + System.out.printf("Inserted %d nodes%n", builder.getGraph().size(0)); + + builder.cleanup(); + graph = builder.getGraph(); + } + + // ── 5. Inspect the result ────────────────────────────────────────────── + try (var view = graph.getView()) { + System.out.printf("Graph built: %d nodes, max level %d, entry node %s%n", + graph.size(0), + graph.getMaxLevel(), + view.entryNode()); + } + System.out.printf("RAM used: %.1f KB%n", graph.ramBytesUsed() / 1024.0); + + // ── 6. Save to disk ─────────────────────────────────────────────────── + // InlineVectors stores float32 versions of the int8 components for on-disk reranking. + Path graphPath = Files.createTempFile("int8-siftsmall", ".jvector"); + try (GraphIndexWriter writer = GraphIndexWriter + .getBuilderFor(GraphIndexWriterTypes.RANDOM_ACCESS_PARALLEL, graph, graphPath) + .with(new InlineVectors(DIM)) + .build()) { + writer.write(Map.of( + FeatureId.INLINE_VECTORS, + nodeId -> { + var bs = ravv.getVector(nodeId); + var fv = vts.createFloatVector(DIM); + for (int i = 0; i < DIM; i++) fv.set(i, bs.get(i)); + return new InlineVectors.State(fv); + } + )); + } + System.out.printf("Graph written to %s (%.1f KB)%n", + graphPath, Files.size(graphPath) / 1024.0); + + // ── 7. Load from disk ───────────────────────────────────────────────── + 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()); + + // ── 8. Search with every query vector ───────────────────────────────── + 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++) { + ByteSequence query = queryVectors.get(q); + var sf = (ScoreFunction.ExactScoreFunction) + node2 -> ByteVectorSimilarityFunction.DOT_PRODUCT.compare(query, ravv.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; change to result.getNodes() for 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()); + } + } + + // ── 9. 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..9cb5ec57d 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 "int8build": + Int8IndexBuild.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/util/SiftLoader.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/SiftLoader.java index a491d0c9e..b5c221143 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/SiftLoader.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/SiftLoader.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.example.util; 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; @@ -81,4 +82,32 @@ public static List> readIvecs(String filename) { return groundTruthTopK; } + + /** + * Reads a .bvecs file (SIFT-style format). + * Each vector is stored as a 4-byte little-endian dimension followed by {@code dim} signed bytes. + * + * @param filePath path to the .bvecs file + * @param dimension expected vector dimensionality (validated against the file's per-vector header) + * @return list of ByteSequence vectors + */ + public static List> readBvecs(String filePath, int dimension) { + var vectors = new ArrayList>(); + try (var dis = new DataInputStream(new BufferedInputStream(new FileInputStream(filePath)))) { + byte[] dimBytes = new byte[4]; + byte[] buf = new byte[dimension]; + while (dis.available() > 0) { + dis.readFully(dimBytes); + int dim = ByteBuffer.wrap(dimBytes).order(ByteOrder.LITTLE_ENDIAN).getInt(); + if (dim != dimension) { + throw new IOException("Expected dimension " + dimension + " but got " + dim); + } + dis.readFully(buf); + vectors.add(vectorTypeSupport.createByteSequence(buf.clone())); + } + } catch (IOException ex) { + throw new UncheckedIOException(ex); + } + return vectors; + } } 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/pom.xml b/jvector-native/pom.xml index ab0090d8d..88073e998 100644 --- a/jvector-native/pom.xml +++ b/jvector-native/pom.xml @@ -141,7 +141,7 @@ ${native.buildtype} false - ${project.basedir}/src/main/native/ + ${project.basedir}/src/main/native/src/ 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_f32.cpp b/jvector-native/src/main/native/benchmarks/bench_similarity_f32.cpp new file mode 100644 index 000000000..b5e31f0cc --- /dev/null +++ b/jvector-native/src/main/native/benchmarks/bench_similarity_f32.cpp @@ -0,0 +1,118 @@ +/* + * 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 fp32 vector similarity kernels: +// cosine_f32, dot_product_f32, euclidean_f32 +// +// 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 float vector: avoids degenerate cosine=NaN cases. +static std::vector make_vec(size_t n, float seed) +{ + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = seed * (1.0f + static_cast(i % 7) * 0.13f); + if (i % 3 == 0) v[i] = -v[i]; + v[i] += 0.5f; + } + return v; +} + +// Benchmark sizes matching production embedding dimensions. +static const std::vector kBenchSizes = {128, 256, 512, 1024, 1536, 3072}; + +// --------------------------------------------------------------------------- +// dot_product_f32 +// --------------------------------------------------------------------------- + +static void BM_dot_product_f32(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_vec(n, 0.7f); + auto b = make_vec(n, 1.3f); + + for (auto _ : state) { + float result = dot_product_f32(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(float)); +} +BENCHMARK(BM_dot_product_f32)->ArgsProduct({kBenchSizes}); + +// --------------------------------------------------------------------------- +// euclidean_f32 +// --------------------------------------------------------------------------- + +static void BM_euclidean_f32(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_vec(n, 0.7f); + auto b = make_vec(n, 1.3f); + + for (auto _ : state) { + float result = euclidean_f32(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(float)); +} +BENCHMARK(BM_euclidean_f32)->ArgsProduct({kBenchSizes}); + +// --------------------------------------------------------------------------- +// cosine_f32 +// --------------------------------------------------------------------------- + +static void BM_cosine_f32(benchmark::State& state) +{ + const size_t n = static_cast(state.range(0)); + auto a = make_vec(n, 0.7f); + auto b = make_vec(n, 1.3f); + + for (auto _ : state) { + float result = cosine_f32(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(float)); +} +BENCHMARK(BM_cosine_f32)->ArgsProduct({kBenchSizes}); + +// --------------------------------------------------------------------------- +// Entry point — benchmark::Initialize parses --benchmark_* flags. +// --------------------------------------------------------------------------- + +BENCHMARK_MAIN(); 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/jvector_avx3_dl_kernels.cpp b/jvector-native/src/main/native/jvector_avx3_dl_kernels.cpp deleted file mode 100644 index ae8ab73bb..000000000 --- a/jvector-native/src/main/native/jvector_avx3_dl_kernels.cpp +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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. - */ - -// AVX3_DL (Ice Lake / Icelake-SP) tier: ONLY kernels that require ICX-specific -// instructions unavailable in the AVX3 (-march=skylake-avx512) compilation -// belong here. Generic kernels that Highway auto-vectorises identically under -// both marches must go in jvector_simd_kernels.cpp instead — that file is -// compiled once for AVX3 and its function pointers are reused by this tier and -// AVX3_SPR via vtable inheritance, avoiding any duplication in .text. -// -// ICX adds over AVX3 (HWY_TARGET_STR_AVX3_DL): -// VNNI, VBMI, VBMI2, IFMA, BITALG, VPOPCNTDQ, GFNI, VAES, VPCLMULQDQ -// -// Compiled with -march=icelake-server. -// Highway will select HWY_AVX3_DL as the static target. -#include "jvector_simd.h" -#include "hwy/highway.h" -#include "assert_hwy_targets.h" - -namespace hn = hwy::HWY_NAMESPACE; - -namespace AVX3_DL { - -} // namespace AVX3_DL diff --git a/jvector-native/src/main/native/meson.build b/jvector-native/src/main/native/meson.build index 9744ddb69..00d29d839 100644 --- a/jvector-native/src/main/native/meson.build +++ b/jvector-native/src/main/native/meson.build @@ -54,7 +54,7 @@ isa_libs = [] foreach isa : isa_variants lib = static_library( 'simdKernels_' + isa['name'], - sources : 'jvector_simd_kernels.cpp', + sources : 'src/jvector_simd_kernels.cpp', include_directories: hwy_inc, cpp_args : isa['args'] + ['-DJV_ISA=' + isa['namespace'], '-fvisibility=hidden'] ) @@ -65,7 +65,7 @@ endforeach # set (AVX3 + VNNI, VBMI, VBMI2, IFMA, BITALG, VPOPCNTDQ, GFNI, VAES, VPCLMULQDQ). avx3_dl_lib = static_library( 'simdKernels_avx3_dl', - sources : 'jvector_avx3_dl_kernels.cpp', + sources : 'src/jvector_avx3_dl_kernels.cpp', include_directories: hwy_inc, cpp_args : [ '-march=icelake-server', @@ -81,7 +81,7 @@ isa_libs += avx3_dl_lib # Requires GCC >= 12 or Clang >= 14. avx3_spr_lib = static_library( 'simdKernels_avx3_spr', - sources : 'jvector_avx3_spr_kernels.cpp', + sources : 'src/jvector_avx3_spr_kernels.cpp', include_directories: hwy_inc, cpp_args : [ '-march=sapphirerapids', @@ -100,9 +100,9 @@ isa_libs += avx3_spr_lib # projects that manage their own dispatch (google/highway#1935). vectorutil_lib = shared_library( 'jvector', - sources : ['jvector_simd.cpp', + sources : ['src/jvector_simd.cpp', 'third_party/highway/hwy/abort.cc'], - include_directories: [include_directories('.'), hwy_inc], + include_directories: [include_directories('src'), hwy_inc], cpp_args : ['-DJVECTOR_BUILD', '-fvisibility=hidden'], link_whole : isa_libs, version : meson.project_version(), @@ -112,7 +112,7 @@ vectorutil_lib = shared_library( # Dependency object for use by executables/tests in this build tree. vectorutil_dep = declare_dependency( link_with : vectorutil_lib, - include_directories: include_directories('.'), + include_directories: include_directories('src'), ) ## Example driver that exercises the runtime-dispatch API. @@ -121,40 +121,37 @@ vectorutil_dep = declare_dependency( # sources : 'examples/cpp_driver.cpp', # dependencies: vectorutil_dep, #) -# -## ---- Tests ----------------------------------------------------------------- -#gtest_dep = dependency('gtest_main', required: true) -# -#test_exe = executable( -# 'test_kernels', -# sources : [ -# 'tests/test_kernels.cpp', -# 'tests/test_cpuFeatures.cpp', -# ], -# dependencies: [vectorutil_dep, gtest_dep], -#) -# -#test('kernels', test_exe, protocol: 'gtest', suite: 'kernels') -#test('cpu_features', test_exe, protocol: 'gtest', suite: 'cpu', -# args: ['--gtest_filter=CpuFeaturesTest.*']) -# -## ---- Benchmarks ------------------------------------------------------------ -#gbench_dep = dependency('benchmark', required: false) -#if gbench_dep.found() -# executable( -# 'bench_kernels', -# sources : 'benchmarks/bench_kernels.cpp', -# dependencies: [vectorutil_dep, gbench_dep], -# cpp_args : ['-O3'], -# ) -#endif -# -#rust_enabled = add_languages('rust', required: false) -#if rust_enabled -# executable( -# 'rust_driver', -# 'examples/rust_driver.rs', -# link_with: vectorutil_lib, -# ) -#endif -# \ No newline at end of file + +# ---- Tests ----------------------------------------------------------------- +gtest_dep = dependency('gtest_main', required: false) + +if gtest_dep.found() + simd_kernels_test = executable( + 'test_simd_kernels', + sources : [ + 'tests/test_helpers.cpp', + 'tests/test_similarity.cpp', + 'tests/test_similarity_i8.cpp', + 'tests/test_elementwise.cpp', + 'tests/test_cpu_features.cpp', + ], + dependencies: [vectorutil_dep, gtest_dep], + ) + + test('simd_kernels', simd_kernels_test, protocol: 'gtest', suite: 'simd_kernels') + +endif +# ---- Benchmarks ------------------------------------------------------------ +gbench_dep = dependency('benchmark', required: false) + +if gbench_dep.found() + executable( + 'bench_simd_kernels', + sources : [ + 'benchmarks/bench_similarity_f32.cpp', + 'benchmarks/bench_similarity_i8.cpp', + ], + dependencies: [vectorutil_dep, gbench_dep], + cpp_args : ['-O3'], + ) +endif \ No newline at end of file diff --git a/jvector-native/src/main/native/assert_hwy_targets.h b/jvector-native/src/main/native/src/assert_hwy_targets.h similarity index 100% rename from jvector-native/src/main/native/assert_hwy_targets.h rename to jvector-native/src/main/native/src/assert_hwy_targets.h diff --git a/jvector-native/src/main/native/jextract_vector_simd.sh b/jvector-native/src/main/native/src/jextract_vector_simd.sh similarity index 74% rename from jvector-native/src/main/native/jextract_vector_simd.sh rename to jvector-native/src/main/native/src/jextract_vector_simd.sh index a2d704c83..be891a7c6 100755 --- a/jvector-native/src/main/native/jextract_vector_simd.sh +++ b/jvector-native/src/main/native/src/jextract_vector_simd.sh @@ -2,6 +2,8 @@ # fail on error set -e +# print commands as they are executed +set +x # Copyright DataStax, Inc. # @@ -17,6 +19,22 @@ set -e # See the License for the specific language governing permissions and # limitations under the License. +# --------------------------------------------------------------------------- +# Path anchors — all derived from the git repository root so the script works +# regardless of the working directory it is invoked from (Maven sets +# workingDirectory to the src directory, but developers may run it from +# anywhere inside the repo). +# --------------------------------------------------------------------------- +REPO_ROOT="$(git -C "$(dirname "$0")" rev-parse --show-toplevel)" +SCRIPT_DIR="${REPO_ROOT}/jvector-native/src/main/native/src" +NATIVE_DIR="${REPO_ROOT}/jvector-native/src/main/native" +MODULE_ROOT="${REPO_ROOT}/jvector-native" + +HIGHWAY_DIR="${NATIVE_DIR}/third_party/highway" +BUILD_DIR="${MODULE_ROOT}/target/meson-build" +RESOURCES_DIR="${MODULE_ROOT}/src/main/resources" +JAVA_OUT_DIR="${MODULE_ROOT}/src/main/java" + if [ "$1" == "--auto-install-deps" ] ; then AUTO_INSTALL_DEPS=true ; shift ; fi printf "AUTO_INSTALL_DEPS=%s\n" "${AUTO_INSTALL_DEPS}" @@ -29,13 +47,13 @@ if [ "$BUILDTYPE" != "release" ] && [ "$BUILDTYPE" != "debug" ] && [ "$BUILDTYPE fi printf "BUILDTYPE=%s\n" "${BUILDTYPE}" -mkdir -p ../resources +mkdir -p "${RESOURCES_DIR}" + # compile jvector_simd_check.cpp as x86-64 # compile jvector_simd.cpp as skylake-avx512 # produce one shared library # Check that the Google Highway submodule has been initialised -HIGHWAY_DIR="third_party/highway" if [ ! -f "${HIGHWAY_DIR}/hwy/highway.h" ]; then echo "ERROR: Google Highway submodule not found at ${HIGHWAY_DIR}." echo " Run the following command from the repository root to fix this:" @@ -80,24 +98,23 @@ if [ "$(printf '%s\n' "$MIN_GCC_VERSION" "$CURRENT_GPP_VERSION" | sort -V | head exit 1 fi -BUILD_DIR="../../../target/meson-build" -rm -rf ../resources/libjvector.so +rm -rf "${RESOURCES_DIR}/libjvector.so" # Configure (--wipe resets any stale configuration) then compile -meson setup "${BUILD_DIR}" \ +meson setup "${BUILD_DIR}" "${NATIVE_DIR}" \ --wipe \ --buildtype="${BUILDTYPE}" meson compile -C "${BUILD_DIR}" # The versioned .so (e.g. libjvector.so.0.1.0) is the real file; symlinks point to it. -# Copy it to ../resources/ as the plain libjvector.so for Java System.load(). +# Copy it to src/main/resources/ so Maven packages it into the jar for LibraryLoader. SOFILE=$(find "${BUILD_DIR}" -maxdepth 1 -name 'libjvector.so.*' -type f | head -1) if [ -z "${SOFILE}" ]; then echo "ERROR: libjvector.so not found in ${BUILD_DIR} after build." exit 1 fi -cp "${SOFILE}" ../resources/libjvector.so +cp "${SOFILE}" "${RESOURCES_DIR}/libjvector.so" # Generate Java source code # Should only be run when c header changes @@ -109,11 +126,12 @@ then fi jextract \ - --output ../java \ + --output "${JAVA_OUT_DIR}" \ -t io.github.jbellis.jvector.vector.cnative \ - -I . \ + -I "${SCRIPT_DIR}" \ --header-class-name NativeSimdOps \ - jvector_simd.h + "${SCRIPT_DIR}/jvector_simd.h" # Set critical linker option with heap-based segments for all generated methods -sed -i 's/DESC)/DESC, Linker.Option.critical(true))/g' ../java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java +sed -i 's/DESC)/DESC, Linker.Option.critical(true))/g' \ + "${JAVA_OUT_DIR}/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java" 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 new file mode 100644 index 000000000..11e7ed33d --- /dev/null +++ b/jvector-native/src/main/native/src/jvector_avx3_dl_kernels.cpp @@ -0,0 +1,382 @@ +/* + * 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. + */ + +// AVX3_DL (Ice Lake / Icelake-SP) tier: ONLY kernels that require ICX-specific +// instructions unavailable in the AVX3 (-march=skylake-avx512) compilation +// belong here. Generic kernels that Highway auto-vectorises identically under +// both marches must go in jvector_simd_kernels.cpp instead — that file is +// compiled once for AVX3 and its function pointers are reused by this tier and +// AVX3_SPR via vtable inheritance, avoiding any duplication in .text. +// +// ICX adds over AVX3 (HWY_TARGET_STR_AVX3_DL): +// VNNI, VBMI, VBMI2, IFMA, BITALG, VPOPCNTDQ, GFNI, VAES, VPCLMULQDQ +// +// Compiled with -march=icelake-server. +// +// 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" + +// ============================================================================= +// 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/jvector_avx3_spr_kernels.cpp b/jvector-native/src/main/native/src/jvector_avx3_spr_kernels.cpp similarity index 100% rename from jvector-native/src/main/native/jvector_avx3_spr_kernels.cpp rename to jvector-native/src/main/native/src/jvector_avx3_spr_kernels.cpp diff --git a/jvector-native/src/main/native/jvector_cpu_features.h b/jvector-native/src/main/native/src/jvector_cpu_features.h similarity index 100% rename from jvector-native/src/main/native/jvector_cpu_features.h rename to jvector-native/src/main/native/src/jvector_cpu_features.h diff --git a/jvector-native/src/main/native/jvector_simd.cpp b/jvector-native/src/main/native/src/jvector_simd.cpp similarity index 96% rename from jvector-native/src/main/native/jvector_simd.cpp rename to jvector-native/src/main/native/src/jvector_simd.cpp index 8bf8ebef5..adc6ba3ac 100644 --- a/jvector-native/src/main/native/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/jvector_simd.h b/jvector-native/src/main/native/src/jvector_simd.h similarity index 100% rename from jvector-native/src/main/native/jvector_simd.h rename to jvector-native/src/main/native/src/jvector_simd.h diff --git a/jvector-native/src/main/native/jvector_simd_kernel_list.h b/jvector-native/src/main/native/src/jvector_simd_kernel_list.h similarity index 90% rename from jvector-native/src/main/native/jvector_simd_kernel_list.h rename to jvector-native/src/main/native/src/jvector_simd_kernel_list.h index 63e7baa4d..99bd99245 100644 --- a/jvector-native/src/main/native/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/jvector_simd_kernels.cpp b/jvector-native/src/main/native/src/jvector_simd_kernels.cpp similarity index 88% rename from jvector-native/src/main/native/jvector_simd_kernels.cpp rename to jvector-native/src/main/native/src/jvector_simd_kernels.cpp index f4e8c2453..1e13dab55 100644 --- a/jvector-native/src/main/native/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/jvector_simd_kernels.h b/jvector-native/src/main/native/src/jvector_simd_kernels.h similarity index 100% rename from jvector-native/src/main/native/jvector_simd_kernels.h rename to jvector-native/src/main/native/src/jvector_simd_kernels.h diff --git a/jvector-native/src/main/native/tests/test_cpu_features.cpp b/jvector-native/src/main/native/tests/test_cpu_features.cpp new file mode 100644 index 000000000..a921660b6 --- /dev/null +++ b/jvector-native/src/main/native/tests/test_cpu_features.cpp @@ -0,0 +1,263 @@ +/* + * 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. + */ + +// Validates that the native dispatcher selects the ISA tier that matches the +// CPU capabilities reported in /proc/cpuinfo, respecting any JVECTOR_MAX_ISA cap. +// +// Logic mirrors DispatcherCpuFlagsTest.java and the C implementation in +// jvector_cpu_features.h / jvector_simd.cpp exactly. +// +// /proc/cpuinfo is the authoritative ground-truth: the kernel only exposes a +// flag when the OS context-switch support (XCR0) is also in place, so checking +// it is equivalent to checking CPUID + XCR0 together. + +#include "test_helpers.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +// --------------------------------------------------------------------------- +// /proc/cpuinfo helpers — mirrors DispatcherCpuFlagsTest.java +// --------------------------------------------------------------------------- + +// Parse the flags line from the first processor entry in /proc/cpuinfo. +// Returns an empty set if unavailable (non-Linux, non-x86, or unreadable). +static std::unordered_set parse_cpuinfo_flags() +{ + std::unordered_set flags; + std::ifstream f("/proc/cpuinfo"); + if (!f.is_open()) return flags; + + std::string line; + while (std::getline(f, line)) { + if (line.rfind("flags", 0) != 0) continue; + auto colon = line.find(':'); + if (colon == std::string::npos) continue; + std::istringstream iss(line.substr(colon + 1)); + std::string token; + while (iss >> token) flags.insert(token); + break; + } + return flags; +} + +// Tier names in ascending capability order — index is ordinal (mirrors Java). +static const std::vector kIsaTiers = { + "sse42", "avx2", "avx3", "avx3_dl", "avx3_spr" +}; + +static int tier_index(const std::string& name) +{ + auto it = std::find(kIsaTiers.begin(), kIsaTiers.end(), name); + return (it == kIsaTiers.end()) ? -1 : static_cast(it - kIsaTiers.begin()); +} + +// ---- Composite tier predicates (flag names from /proc/cpuinfo) ---- +// These must stay in sync with DispatcherCpuFlagsTest.java and +// jvector_cpu_features.h. The kernel's naming is inconsistent: +// no underscore: avx512f bw cd dq vl, avx512vbmi, avx512ifma +// with underscore: avx512_vnni, avx512_vbmi2, avx512_bitalg, +// avx512_vpopcntdq, avx512_fp16 +// gfni / vaes / vpclmulqdq have no avx512 prefix at all. + +static bool has_avx3(const std::unordered_set& f) +{ + return f.count("avx512f") && f.count("avx512bw") + && f.count("avx512cd") && f.count("avx512dq") + && f.count("avx512vl"); +} + +static bool has_avx3_dl(const std::unordered_set& f) +{ + return has_avx3(f) + && f.count("avx512_vnni") && f.count("avx512vbmi") + && f.count("avx512_vbmi2") && f.count("avx512ifma") + && f.count("avx512_bitalg") && f.count("avx512_vpopcntdq") + && f.count("gfni") && f.count("vaes") + && f.count("vpclmulqdq"); +} + +static bool has_avx3_spr(const std::unordered_set& f) +{ + return has_avx3_dl(f) && f.count("avx512_fp16"); +} + +// Compute the expected ISA tier from /proc/cpuinfo flags and the cap, +// mirroring expectedIsaFromCpuInfo() in DispatcherCpuFlagsTest.java. +static std::string expected_isa(const std::unordered_set& flags, + const std::string& cap) +{ + std::string best; + if (has_avx3_spr(flags)) best = "avx3_spr"; + else if (has_avx3_dl(flags)) best = "avx3_dl"; + else if (has_avx3(flags)) best = "avx3"; + else if (flags.count("avx2")) best = "avx2"; + else best = "sse42"; + + // Clamp down to cap if set and below best. + if (!cap.empty() && tier_index(cap) < tier_index(best)) + return cap; + return best; +} + +// --------------------------------------------------------------------------- +// Fixture — state shared across all tests +// --------------------------------------------------------------------------- + +class CpuFeaturesTest : public ::testing::Test +{ +protected: + static void SetUpTestSuite() + { + s_flags = parse_cpuinfo_flags(); + + const char* active = jvector_simd_get_active_isa(); + const char* cap_c = jvector_simd_get_max_isa_env(); + s_active = active ? active : ""; + s_cap = cap_c ? cap_c : ""; + + // Detect CPU emulators (e.g. Intel SDE): they intercept CPUID and + // return synthetic features, but /proc/cpuinfo still reflects the host. + // When the active ISA cannot be explained by the host's cpuinfo flags + // the comparison tests are meaningless, so we skip them. + std::string host_expected = expected_isa(s_flags, s_cap); + s_under_emulator = !s_flags.empty() + && tier_index(s_active) > tier_index(host_expected); + + s_available = !s_flags.empty() && !s_under_emulator; + + std::printf("[ CPU ] active_isa=%s JVECTOR_MAX_ISA=%s " + "cpuinfo_flags=%zu emulator=%s\n", + s_active.c_str(), + s_cap.empty() ? "(unset)" : s_cap.c_str(), + s_flags.size(), + s_under_emulator ? "yes (cpuinfo skipped)" : "no"); + } + + static bool isCappedBelow(const std::string& tier) + { + return !s_cap.empty() && tier_index(s_cap) < tier_index(tier); + } + + static std::unordered_set s_flags; + static std::string s_active; + static std::string s_cap; + static bool s_available; + static bool s_under_emulator; +}; + +std::unordered_set CpuFeaturesTest::s_flags; +std::string CpuFeaturesTest::s_active; +std::string CpuFeaturesTest::s_cap; +bool CpuFeaturesTest::s_available = false; +bool CpuFeaturesTest::s_under_emulator = false; + +#define SKIP_IF_UNAVAILABLE() \ + do { \ + if (s_flags.empty()) GTEST_SKIP() << "/proc/cpuinfo unavailable"; \ + if (s_under_emulator) GTEST_SKIP() << "CPU emulator detected (SDE?): " \ + "active_isa=" << s_active << " exceeds host cpuinfo capability"; \ + } while (0) + +// --------------------------------------------------------------------------- +// Tests — mirror each @Test method in DispatcherCpuFlagsTest.java +// --------------------------------------------------------------------------- + +// AVX2 tier is selected when avx2 is present and the cap allows it. +TEST_F(CpuFeaturesTest, Avx2Detection) +{ + SKIP_IF_UNAVAILABLE(); + bool cpu_has_avx2 = s_flags.count("avx2") > 0; + + if (cpu_has_avx2 && !isCappedBelow("avx2")) { + EXPECT_GE(tier_index(s_active), tier_index("avx2")) + << "Expected AVX2 or higher when avx2 flag present, got: " << s_active; + } else if (!cpu_has_avx2) { + EXPECT_EQ(s_active, "sse42") + << "Expected sse42 when avx2 flag absent, got: " << s_active; + } +} + +// AVX3 tier is selected when all avx512 baseline flags are present. +TEST_F(CpuFeaturesTest, Avx3Detection) +{ + SKIP_IF_UNAVAILABLE(); + bool cpu_has_avx3 = has_avx3(s_flags); + + if (cpu_has_avx3 && !isCappedBelow("avx3")) { + EXPECT_GE(tier_index(s_active), tier_index("avx3")) + << "Expected AVX3 or higher when avx512 baseline flags present, got: " << s_active; + } else if (!cpu_has_avx3) { + EXPECT_LT(tier_index(s_active), tier_index("avx3")) + << "Expected below AVX3 when avx512 baseline flags absent, got: " << s_active; + } +} + +// AVX3_DL tier is selected when all ICX flags are present on top of AVX3. +TEST_F(CpuFeaturesTest, Avx3DlDetection) +{ + SKIP_IF_UNAVAILABLE(); + bool cpu_has_avx3_dl = has_avx3_dl(s_flags); + + if (cpu_has_avx3_dl && !isCappedBelow("avx3_dl")) { + EXPECT_GE(tier_index(s_active), tier_index("avx3_dl")) + << "Expected AVX3_DL or higher when all ICX flags present, got: " << s_active; + } else if (!cpu_has_avx3_dl) { + EXPECT_LT(tier_index(s_active), tier_index("avx3_dl")) + << "Expected below AVX3_DL when ICX flags absent, got: " << s_active; + } +} + +// AVX3_SPR tier is selected when avx512_fp16 and all ICX flags are present. +TEST_F(CpuFeaturesTest, Avx3SprDetection) +{ + SKIP_IF_UNAVAILABLE(); + bool cpu_has_avx3_spr = has_avx3_spr(s_flags); + + if (cpu_has_avx3_spr && !isCappedBelow("avx3_spr")) { + EXPECT_EQ(s_active, "avx3_spr") + << "Expected avx3_spr when fp16 + all ICX flags present and uncapped"; + } else if (!cpu_has_avx3_spr) { + EXPECT_LT(tier_index(s_active), tier_index("avx3_spr")) + << "Expected below AVX3_SPR when avx512_fp16 absent, got: " << s_active; + } +} + +// End-to-end: the tier the dispatcher chose must match what /proc/cpuinfo implies. +TEST_F(CpuFeaturesTest, DispatcherMatchesCpuInfo) +{ + SKIP_IF_UNAVAILABLE(); + + std::string exp = expected_isa(s_flags, s_cap); + + // Collect all flags for the failure message. + std::string all_flags = std::accumulate( + s_flags.begin(), s_flags.end(), std::string{}, + [](const std::string& a, const std::string& b) { + return a.empty() ? b : a + " " + b; + }); + + EXPECT_EQ(s_active, exp) + << "Dispatcher chose '" << s_active + << "' but /proc/cpuinfo implies '" << exp << "'." + << "\nCPU flags: " << all_flags; +} diff --git a/jvector-native/src/main/native/tests/test_elementwise.cpp b/jvector-native/src/main/native/tests/test_elementwise.cpp new file mode 100644 index 000000000..04c836461 --- /dev/null +++ b/jvector-native/src/main/native/tests/test_elementwise.cpp @@ -0,0 +1,193 @@ +/* + * 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 element-wise in-place arithmetic kernels: +// add_in_place_f32, add_scalar_in_place_f32, +// sub_in_place_f32, sub_scalar_in_place_f32, +// max_f32, min_in_place_f32. +// +// Vector sizes cover the same ISA boundary / tail matrix as test_similarity.cpp: +// size 1 — tail only +// size 3 — tail only +// size 4 — SSE42 exact / AVX2+AVX512 capped path +// size 7 — SSE42 1 full + 3-tail +// size 8 — AVX2 exact / AVX512 capped +// size 15 — AVX2 1 full + 7-tail +// size 16 — SSE42 4× main exact / AVX512 1 full +// size 17 — SSE42 4× main + 1-tail +// size 32 — AVX2 4× main exact +// size 37 — AVX2 4× main + 5-tail +// size 64 — AVX512 4× main exact +// size 71 — AVX512 4× main + 7-tail +// size 100, 128, 255 — large mixed / power-of-2 / odd + +#include "test_helpers.h" + +// --------------------------------------------------------------------------- +// Fixture +// --------------------------------------------------------------------------- + +class ElementWiseTest : public ::testing::TestWithParam {}; + +// --------------------------------------------------------------------------- +// add_in_place_f32: v1[i] += v2[i] +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, AddInPlace) +{ + const size_t n = GetParam().length; + auto v1 = make_vec(n, 1.1f); + auto v2 = make_vec(n, 0.7f); + auto want = v1; + for (size_t i = 0; i < n; ++i) want[i] += v2[i]; + + auto got = v1; + add_in_place_f32(got.data(), v2.data(), n); + + for (size_t i = 0; i < n; ++i) + EXPECT_NEAR(got[i], want[i], 1e-5f) << "add_in_place_f32[" << i << "]"; +} + +// --------------------------------------------------------------------------- +// add_scalar_in_place_f32: v1[i] += scalar +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, AddScalarInPlace) +{ + const size_t n = GetParam().length; + const float scalar = 3.14f; + auto v1 = make_vec(n, 1.1f); + auto want = v1; + for (size_t i = 0; i < n; ++i) want[i] += scalar; + + auto got = v1; + add_scalar_in_place_f32(got.data(), scalar, n); + + for (size_t i = 0; i < n; ++i) + EXPECT_NEAR(got[i], want[i], 1e-5f) << "add_scalar_in_place_f32[" << i << "]"; +} + +// --------------------------------------------------------------------------- +// sub_in_place_f32: v1[i] -= v2[i] +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, SubInPlace) +{ + const size_t n = GetParam().length; + auto v1 = make_vec(n, 1.1f); + auto v2 = make_vec(n, 0.7f); + auto want = v1; + for (size_t i = 0; i < n; ++i) want[i] -= v2[i]; + + auto got = v1; + sub_in_place_f32(got.data(), v2.data(), n); + + for (size_t i = 0; i < n; ++i) + EXPECT_NEAR(got[i], want[i], 1e-5f) << "sub_in_place_f32[" << i << "]"; +} + +// --------------------------------------------------------------------------- +// sub_scalar_in_place_f32: v1[i] -= scalar +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, SubScalarInPlace) +{ + const size_t n = GetParam().length; + const float scalar = 2.71f; + auto v1 = make_vec(n, 1.1f); + auto want = v1; + for (size_t i = 0; i < n; ++i) want[i] -= scalar; + + auto got = v1; + sub_scalar_in_place_f32(got.data(), scalar, n); + + for (size_t i = 0; i < n; ++i) + EXPECT_NEAR(got[i], want[i], 1e-5f) << "sub_scalar_in_place_f32[" << i << "]"; +} + +// --------------------------------------------------------------------------- +// max_f32: returns the maximum element +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, MaxF32) +{ + const size_t n = GetParam().length; + auto v = make_vec(n, 0.9f); + + float want = *std::max_element(v.begin(), v.end()); + float got = max_f32(v.data(), n); + + EXPECT_FLOAT_EQ(got, want); +} + +// max_f32 on a vector with a known maximum at the last position (tail element) +TEST_P(ElementWiseTest, MaxF32TailElement) +{ + const size_t n = GetParam().length; + auto v = make_vec(n, 0.5f); + // Place the global maximum in the very last element — exercises tail path. + v.back() = 1e6f; + + float got = max_f32(v.data(), n); + + EXPECT_FLOAT_EQ(got, 1e6f); +} + +// --------------------------------------------------------------------------- +// min_in_place_f32: v1[i] = min(v1[i], v2[i]) +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, MinInPlace) +{ + const size_t n = GetParam().length; + auto v1 = make_vec(n, 1.1f); + auto v2 = make_vec(n, 0.7f); + auto want = v1; + for (size_t i = 0; i < n; ++i) want[i] = std::min(want[i], v2[i]); + + auto got = v1; + min_in_place_f32(got.data(), v2.data(), n); + + for (size_t i = 0; i < n; ++i) + EXPECT_NEAR(got[i], want[i], 1e-5f) << "min_in_place_f32[" << i << "]"; +} + +// add then sub back — result must equal the original vector +TEST_P(ElementWiseTest, AddSubRoundTrip) +{ + const size_t n = GetParam().length; + auto original = make_vec(n, 1.3f); + auto delta = make_vec(n, 0.4f); + + auto v = original; + add_in_place_f32(v.data(), delta.data(), n); + sub_in_place_f32(v.data(), delta.data(), n); + + for (size_t i = 0; i < n; ++i) + EXPECT_NEAR(v[i], original[i], 1e-5f) << "add_sub_roundtrip[" << i << "]"; +} + +// --------------------------------------------------------------------------- +// Instantiation +// --------------------------------------------------------------------------- + +INSTANTIATE_TEST_SUITE_P( + AllSizes, + ElementWiseTest, + ::testing::ValuesIn(kKernelTestParams), + [](const ::testing::TestParamInfo& info) { + return info.param.description; + }); diff --git a/jvector-native/src/main/native/tests/test_helpers.cpp b/jvector-native/src/main/native/tests/test_helpers.cpp new file mode 100644 index 000000000..45b35cc6e --- /dev/null +++ b/jvector-native/src/main/native/tests/test_helpers.cpp @@ -0,0 +1,110 @@ +/* + * 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. + */ + +#include "test_helpers.h" + +// --------------------------------------------------------------------------- +// Global test environment — prints the active ISA once for the whole binary. +// Registered via AddGlobalTestEnvironment at static-init time so it fires +// before any test suite runs, regardless of which .cpp files are linked. +// --------------------------------------------------------------------------- + +class JVectorIsaEnvironment : public ::testing::Environment +{ +public: + void SetUp() override + { + std::printf("[ ISA ] Active dispatch tier: %s\n", + jvector_simd_get_active_isa()); + } +}; + +static ::testing::Environment* const kIsaEnv = + ::testing::AddGlobalTestEnvironment(new JVectorIsaEnvironment); + +// --------------------------------------------------------------------------- +// make_vec: deterministic test vectors. +// Values in roughly (-2, 2] with a mix of signs so no element is zero. +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// Canonical test sizes — shared by all parametrised suites. +// Covers tail-only, single-register, and 4x-unrolled main-loop paths for +// SSE42 (4 lanes), AVX2 (8 lanes), and AVX512 (16 lanes). +// --------------------------------------------------------------------------- + +const std::vector kKernelTestParams = { + // ---- tail-only (< 4 lanes for any ISA) -------------------------------- + { 1, "tail_1_all_isa"}, + { 3, "tail_3_all_isa"}, + // ---- SSE42 boundary (4-lane register) ---------------------------------- + { 4, "sse42_exact_4"}, + { 5, "sse42_1full_tail_1"}, + { 7, "sse42_1full_tail_3"}, + // ---- AVX2 boundary (8-lane register) ----------------------------------- + { 8, "avx2_exact_8"}, + { 9, "avx2_1full_tail_1"}, + { 15, "avx2_1full_tail_7"}, + // ---- SSE42 4x-unrolled main loop (16 elements = 4 × 4 lanes) ---------- + { 16, "sse42_4x_main_exact"}, + { 17, "sse42_4x_main_tail_1"}, + { 19, "sse42_4x_main_tail_3"}, + // ---- AVX2 4x-unrolled main loop (32 elements = 4 × 8 lanes) ----------- + { 32, "avx2_4x_main_exact"}, + { 33, "avx2_4x_main_tail_1"}, + { 37, "avx2_4x_main_tail_5"}, + // ---- AVX512 boundary (16-lane register) -------------------------------- + { 64, "avx512_4x_main_exact"}, + { 71, "avx512_4x_main_tail_7"}, + // ---- Odd large size exercising all loop stages ------------------------- + {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) +{ + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + v[i] = seed * (1.0f + static_cast(i % 7) * 0.13f); + if (i % 3 == 0) v[i] = -v[i]; // mix of signs + v[i] += 0.5f; // ensure non-zero even after sign flip + } + 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 new file mode 100644 index 000000000..e1e9cd79c --- /dev/null +++ b/jvector-native/src/main/native/tests/test_helpers.h @@ -0,0 +1,57 @@ +/* + * 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. + */ + +// Shared helpers for the test_simd_kernels test binary. +// Included by each test .cpp file; defined in test_helpers.cpp. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "jvector_simd.h" + +// --------------------------------------------------------------------------- +// Deterministic test vectors. +// make_vec(n, seed) produces n floats with a mix of signs and magnitudes +// so that no element is exactly zero (important for cosine tests). +// --------------------------------------------------------------------------- + +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. +// Used by every parametrised test suite in the binary so the same set of +// sizes exercises each kernel. +// --------------------------------------------------------------------------- + +struct KernelTestParam { + size_t length; + std::string description; +}; + +// The canonical set of sizes that hits every code path across ISA tiers. +// See the top-of-file comment in test_similarity.cpp for the full breakdown. +extern const std::vector kKernelTestParams; diff --git a/jvector-native/src/main/native/tests/test_similarity.cpp b/jvector-native/src/main/native/tests/test_similarity.cpp new file mode 100644 index 000000000..c08947e4f --- /dev/null +++ b/jvector-native/src/main/native/tests/test_similarity.cpp @@ -0,0 +1,265 @@ +/* + * 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 vector similarity kernels: cosine_f32, dot_product_f32, euclidean_f32. +// +// The library dispatches to the best ISA available on the host CPU at +// static-init time. Vector sizes are chosen to hit every code path in the +// kernel loops regardless of which ISA is selected: +// +// SSE42 (4 lanes): +// sizes 1, 3 — tail only (< 4) +// size 4 — exact one vector, no tail +// size 7 — one full + 3-element tail +// size 16 — 4x unrolled main loop, no tail +// size 19 — 4x main + 3-element tail +// +// AVX2 (8 lanes): +// sizes 1, 3 — capped fast path (≤4), tail only +// size 4 — capped fast path (≤4), one vector no tail +// size 7 — capped fast path (≤8), tail = 7 < 8 +// size 8 — capped fast path (≤8), exact no tail +// size 15 — one full + 7-element tail +// size 32 — 4x unrolled main loop, no tail +// size 37 — 4x main + 5-element tail +// +// AVX3/AVX512 (16 lanes): +// sizes 1, 3 — capped (≤4), tail only +// size 4 — capped (≤4), exact +// size 8 — capped (≤8), exact +// size 15 — one full (16 lanes) – 1 = tail +// size 16 — exact one full register +// size 64 — 4x unrolled, no tail +// size 71 — 4x main + 7-element tail + +#include "test_helpers.h" + +// --------------------------------------------------------------------------- +// Reference scalar implementations +// --------------------------------------------------------------------------- + +static float ref_dot(const std::vector& a, const std::vector& b) +{ + float s = 0.0f; + for (size_t i = 0; i < a.size(); ++i) s += a[i] * b[i]; + return s; +} + +static float ref_euclidean(const std::vector& a, const std::vector& b) +{ + float s = 0.0f; + for (size_t i = 0; i < a.size(); ++i) { + float d = a[i] - b[i]; + s += d * d; + } + return s; +} + +static float ref_cosine(const std::vector& a, const std::vector& b) +{ + float ab = 0.0f, aa = 0.0f, bb = 0.0f; + for (size_t i = 0; i < a.size(); ++i) { + ab += a[i] * b[i]; + aa += a[i] * a[i]; + bb += b[i] * b[i]; + } + return ab / std::sqrt(aa * bb); +} + +// --------------------------------------------------------------------------- +// Parametrised test fixture +// --------------------------------------------------------------------------- + +class SimilarityTest : public ::testing::TestWithParam {}; + +// --------------------------------------------------------------------------- +// dot_product_f32 — SIMD result must match the scalar reference +// --------------------------------------------------------------------------- + +TEST_P(SimilarityTest, DotProduct) +{ + const size_t n = GetParam().length; + auto a = make_vec(n, 0.7f); + auto b = make_vec(n, 1.3f); + + float want = ref_dot(a, b); + float got = dot_product_f32(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, want, 1e-4f * std::abs(want)); +} + +// --------------------------------------------------------------------------- +// dot_product_f32 with non-zero offsets — exercises the aoffset/boffset path +// --------------------------------------------------------------------------- + +TEST_P(SimilarityTest, DotProductWithOffset) +{ + const size_t n = GetParam().length; + const size_t prefix = 3; // arbitrary prefix that must be ignored + + // Pad the front with values that must not contribute to the result. + std::vector a_pad(prefix + n); + std::vector b_pad(prefix + n); + auto a = make_vec(n, 0.7f); + auto b = make_vec(n, 1.3f); + std::copy(a.begin(), a.end(), a_pad.begin() + prefix); + std::copy(b.begin(), b.end(), b_pad.begin() + prefix); + + float want = ref_dot(a, b); + float got = dot_product_f32(a_pad.data(), prefix, b_pad.data(), prefix, n); + + EXPECT_NEAR(got, want, 1e-4f * std::abs(want)); +} + +// --------------------------------------------------------------------------- +// euclidean_f32 — squared L2 distance +// --------------------------------------------------------------------------- + +TEST_P(SimilarityTest, Euclidean) +{ + const size_t n = GetParam().length; + auto a = make_vec(n, 0.7f); + auto b = make_vec(n, 1.3f); + + float want = ref_euclidean(a, b); + float got = euclidean_f32(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, want, 1e-4f * std::abs(want)); +} + +// --------------------------------------------------------------------------- +// euclidean_f32 — identical vectors should give exactly 0.0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityTest, EuclideanSameVector) +{ + const size_t n = GetParam().length; + auto a = make_vec(n, 0.9f); + + float got = euclidean_f32(a.data(), 0, a.data(), 0, n); + + // Exact zero is expected since a == b; scale tolerance with length to + // allow for FMA reassociation differences across ISAs. + EXPECT_NEAR(got, 0.0f, 1e-6f * static_cast(n)); +} + +// --------------------------------------------------------------------------- +// cosine_f32 — cosine similarity +// --------------------------------------------------------------------------- + +TEST_P(SimilarityTest, Cosine) +{ + const size_t n = GetParam().length; + auto a = make_vec(n, 0.7f); + auto b = make_vec(n, 1.3f); + + float want = ref_cosine(a, b); + float got = cosine_f32(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, want, 1e-4f * std::abs(want)); +} + +// --------------------------------------------------------------------------- +// cosine_f32 — parallel vectors should give similarity = 1.0 +// --------------------------------------------------------------------------- + +TEST_P(SimilarityTest, CosineParallelVectors) +{ + const size_t n = GetParam().length; + auto a = make_vec(n, 1.0f); + + // b = 2*a — same direction, different magnitude → cosine = 1.0 + std::vector b(n); + for (size_t i = 0; i < n; ++i) b[i] = 2.0f * a[i]; + + float got = cosine_f32(a.data(), 0, b.data(), 0, n); + + EXPECT_NEAR(got, 1.0f, 1e-5f); +} + +// --------------------------------------------------------------------------- +// cosine_f32 — orthogonal vectors should give similarity ≈ 0.0 +// +// Orthogonality is constructed analytically for even n (alternating +/-): +// a = [+1, +1, +1, ...] +// b = [+1, -1, +1, -1, ...] — then a·b = 0 if n is even. +// For odd n we only use n-1 elements (prefix) so the dot is still zero. +// --------------------------------------------------------------------------- + +TEST_P(SimilarityTest, 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); // largest even prefix + + std::vector a(n, 0.0f), b(n, 0.0f); + for (size_t i = 0; i < even_n; ++i) { + a[i] = 1.0f; + b[i] = (i % 2 == 0) ? 1.0f : -1.0f; + } + + float got = cosine_f32(a.data(), 0, b.data(), 0, n); + + // Generous tolerance: FP accumulation order differs between ISA tiers. + EXPECT_NEAR(got, 0.0f, 1e-4f); +} + +// --------------------------------------------------------------------------- +// Instantiation — named using the description field +// --------------------------------------------------------------------------- + +INSTANTIATE_TEST_SUITE_P( + AllSizes, + SimilarityTest, + ::testing::ValuesIn(kKernelTestParams), + [](const ::testing::TestParamInfo& info) { + return info.param.description; + }); + +// --------------------------------------------------------------------------- +// ISA-tier sanity test: confirm JVECTOR_MAX_ISA cap is respected when set +// --------------------------------------------------------------------------- + +TEST(IsaDispatch, MaxIsaEnvHonoured) +{ + const char* env = jvector_simd_get_max_isa_env(); + const char* active = jvector_simd_get_active_isa(); + + if (env == nullptr) { + // No override — just report which tier was auto-detected. + SUCCEED() << "No JVECTOR_MAX_ISA set; auto-selected: " << active; + return; + } + + // Tiers ordered by capability (ascending index = lower capability). + static const char* kOrder[] = {"sse42", "avx2", "avx3", "avx3_dl", "avx3_spr"}; + auto tier_idx = [](const char* name) -> int { + for (int i = 0; i < 5; ++i) + if (std::strcmp(kOrder[i], name) == 0) return i; + return -1; + }; + + int env_idx = tier_idx(env); + int active_idx = tier_idx(active); + + ASSERT_GE(env_idx, 0) << "Unrecognised JVECTOR_MAX_ISA value: " << env; + ASSERT_GE(active_idx, 0) << "Unrecognised active ISA: " << active; + + // Active tier must be <= requested cap. + EXPECT_LE(active_idx, env_idx) + << "Active ISA (" << active << ") exceeds requested cap (" << env << ")"; +} 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 4a3b69e93..1a975a487 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 @@ -18,6 +18,7 @@ import com.carrotsearch.randomizedtesting.RandomizedTest; import io.github.jbellis.jvector.TestUtil; +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; @@ -52,6 +53,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); diff --git a/rat-excludes.txt b/rat-excludes.txt index bcdf05ae4..4d0eb0740 100644 --- a/rat-excludes.txt +++ b/rat-excludes.txt @@ -37,3 +37,5 @@ local_datasets/** **/datasets/** jvector-native/src/main/native/third_party/** src/main/native/third_party/** +jvector-native/src/target/meson-build/** +jvector-native/target/meson-build/** diff --git a/siftsmall/siftsmall_base.bvecs b/siftsmall/siftsmall_base.bvecs new file mode 100644 index 000000000..af90201bf Binary files /dev/null and b/siftsmall/siftsmall_base.bvecs differ diff --git a/siftsmall/siftsmall_query.bvecs b/siftsmall/siftsmall_query.bvecs new file mode 100644 index 000000000..9ed28ac65 Binary files /dev/null and b/siftsmall/siftsmall_query.bvecs differ