Date: Mon, 22 Dec 2025 21:33:34 -0800
Subject: [PATCH 001/243] Initial commit of the ASH indexing and vector
aspects, plus some basic testing programs. Not executable yet.
---
.../jvector/quantization/ASHVectors.java | 247 +++++
.../quantization/AsymmetricHashing.java | 937 ++++++++++++++++++
.../jbellis/jvector/example/DistancesASH.java | 153 +++
.../example/TestAshRunWithoutTraining.java | 43 +
4 files changed, 1380 insertions(+)
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
create mode 100644 jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
create mode 100644 jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
new file mode 100644
index 000000000..b1f822ea2
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -0,0 +1,247 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.quantization;
+
+import io.github.jbellis.jvector.disk.RandomAccessReader;
+import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.util.RamUsageEstimator;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.io.DataOutput;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+
+/**
+ * Container for ASH-compressed vectors with a minimal scoring interface.
+ *
+ *
+ * This class intentionally implements only the functionality required for:
+ *
+ * - Encoding benchmarks
+ * - Approximate distance microbenchmarks
+ * - Correctness validation against full-precision distances
+ *
+ *
+ *
+ * TODO (future work):
+ *
+ * - Implement landmark-aware ASH scoring using residualNorm and dotWithLandmark
+ * - Add a full ASH ScoreFunction aligned with graph search semantics
+ * - Support similarityToNeighbor(...) where appropriate
+ * - Integrate ASH as a Feature for on-disk graph indexes
+ * - Define reranking strategy (likely via full vectors)
+ * - Support multi-landmark ASH (this is intentionally deferred)
+ *
+ *
+ *
+ * NOTE:
+ * This class avoids PQ/NVQ feature plumbing and graph integration by design.
+ * Those concerns belong to the scoring plane and will be addressed after
+ * baseline performance is established.
+ */
+public class ASHVectors implements CompressedVectors {
+
+ final AsymmetricHashing ash;
+ final AsymmetricHashing.QuantizedVector[] compressedVectors;
+
+ /**
+ * Initialize ASHVectors with an array of ASH-compressed vectors.
+ *
+ *
+ * The array is treated as immutable after construction.
+ */
+ public ASHVectors(AsymmetricHashing ash,
+ AsymmetricHashing.QuantizedVector[] compressedVectors) {
+ this.ash = ash;
+ this.compressedVectors = compressedVectors;
+ }
+
+ @Override
+ public int count() {
+ return compressedVectors.length;
+ }
+
+ /**
+ * Serialize the compressor followed by the compressed vectors.
+ *
+ *
+ * NOTE:
+ * This format is intentionally simple and versioned only at the compressor level.
+ * Additional per-vector metadata (e.g., multiple landmarks) must be versioned
+ * carefully when introduced.
+ */
+ @Override
+ public void write(DataOutput out, int version) throws IOException {
+ // Write ASH compressor first
+ ash.write(out, version);
+
+ // Write vector count
+ out.writeInt(compressedVectors.length);
+
+ // Write vectors
+ for (var v : compressedVectors) {
+ v.write(out, ash.quantizedDim);
+ }
+ }
+
+ public static ASHVectors load(RandomAccessReader in) throws IOException {
+ var ash = AsymmetricHashing.load(in);
+
+ int size = in.readInt();
+ if (size < 0) {
+ throw new IOException("Invalid compressed vector count " + size);
+ }
+
+ var compressedVectors = new AsymmetricHashing.QuantizedVector[size];
+ for (int i = 0; i < size; i++) {
+ compressedVectors[i] =
+ AsymmetricHashing.QuantizedVector.load(in, ash.quantizedDim);
+ }
+
+ return new ASHVectors(ash, compressedVectors);
+ }
+
+ /**
+ * Return an approximate score function for the given query.
+ *
+ *
+ * Baseline implementation:
+ *
+ * - Encode query once using ASH
+ * - Score candidates using pure Hamming distance
+ *
+ *
+ *
+ * TODO:
+ *
+ * - Replace pure Hamming with full ASH scoring
+ * - Incorporate landmark scalars (dotWithLandmark, residualNorm)
+ * - Respect VectorSimilarityFunction semantics where applicable
+ *
+ */
+ @Override
+ public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(
+ VectorFloat> query,
+ VectorSimilarityFunction similarityFunction) {
+
+ // Encode query once
+ var qv = AsymmetricHashing.QuantizedVector.createEmpty(ash.quantizedDim);
+ ash.encodeTo(query, qv);
+ final long[] queryBits = qv.binaryVector;
+
+ return node -> {
+ var v = compressedVectors[node];
+
+ // TODO:
+ // - Replace with proper ASH scoring.
+ // - Consider affine scaling or normalization.
+ int hamming = VectorUtil.hammingDistance(queryBits, v.binaryVector);
+
+ // Higher similarity is better
+ return -hamming;
+ };
+ }
+
+ /**
+ * For ASH, precomputed and non-precomputed scoring are currently identical.
+ *
+ *
+ * TODO:
+ * - Evaluate caching projected queries or other query-dependent state.
+ */
+ @Override
+ public ScoreFunction.ApproximateScoreFunction precomputedScoreFunctionFor(
+ VectorFloat> query,
+ VectorSimilarityFunction similarityFunction) {
+ return scoreFunctionFor(query, similarityFunction);
+ }
+
+ /**
+ * Diversity-aware scoring is not supported for ASH at this stage.
+ *
+ *
+ * TODO:
+ * - Define whether diversity should be measured via:
+ * * Hamming distance
+ * * Landmark-aware distance
+ * * Full-vector reranking
+ */
+ @Override
+ public ScoreFunction.ApproximateScoreFunction diversityFunctionFor(
+ int node1,
+ VectorSimilarityFunction similarityFunction) {
+ throw new UnsupportedOperationException("ASH diversity scoring not implemented");
+ }
+
+ public AsymmetricHashing.QuantizedVector get(int ordinal) {
+ return compressedVectors[ordinal];
+ }
+
+ @Override
+ public int getOriginalSize() {
+ return ash.originalDimension * Float.BYTES;
+ }
+
+ @Override
+ public int getCompressedSize() {
+ return ash.compressedVectorSize();
+ }
+
+ @Override
+ public AsymmetricHashing getCompressor() {
+ return ash;
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ int REF_BYTES = RamUsageEstimator.NUM_BYTES_OBJECT_REF;
+ int AH_BYTES = RamUsageEstimator.NUM_BYTES_ARRAY_HEADER;
+
+ long vectorsArraySize =
+ AH_BYTES + (long) REF_BYTES * compressedVectors.length;
+ long vectorsDataSize =
+ (long) ash.compressedVectorSize() * compressedVectors.length;
+
+ return ash.ramBytesUsed() + vectorsArraySize + vectorsDataSize;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+
+ ASHVectors that = (ASHVectors) o;
+ return Objects.equals(ash, that.ash)
+ && Arrays.equals(compressedVectors, that.compressedVectors);
+ }
+
+ @Override
+ public int hashCode() {
+ int result = Objects.hash(ash);
+ result = 31 * result + Arrays.hashCode(compressedVectors);
+ return result;
+ }
+
+ @Override
+ public String toString() {
+ return "ASHVectors{count=" + compressedVectors.length +
+ ", ash=" + ash + '}';
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
new file mode 100644
index 000000000..22e1201cf
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -0,0 +1,937 @@
+/*
+ * 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.annotations.VisibleForTesting;
+import io.github.jbellis.jvector.disk.RandomAccessReader;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.util.Accountable;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.ByteSequence;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+
+import java.io.DataOutput;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Objects;
+import java.util.concurrent.ForkJoinPool;
+import java.util.Random;
+import java.util.stream.IntStream;
+import org.apache.commons.math3.linear.*;
+
+import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED;
+
+/**
+ * Asymmetric Hashing (ASH) for float vectors.
+ * Encodes each vector into a fixed-length binary code using a learned or random
+ * orthonormal projection and sign thresholding.
+ */
+public class AsymmetricHashing implements VectorCompressor, Accountable {
+ public static final int ITQ = 1, LANDING = 2, RANDOM = 3;
+ private static final int MAGIC = 0x75EC4012; // TODO update the magic number?
+
+ private static final int HEADER_BITS = 48; // fixed ASH per-vector header
+
+ private static final VectorTypeSupport vectorTypeSupport =
+ VectorizationProvider.getInstance().getVectorTypeSupport();
+
+ // ---------------------------------------------------------------------
+ // Index-wide immutable state
+ // ---------------------------------------------------------------------
+
+ /** Mean of the dataset; also used as degenerate landmark */
+ public final VectorFloat> globalMean;
+
+ /** Original (uncompressed) dimensionality */
+ public final int originalDimension;
+
+ /** Total bits per encoded vector (header + body) */
+ public final int encodedBits;
+
+ /** Number of bits in the binary body */
+ public final int quantizedDim;
+
+ /** Optimizer / learning mode */
+ public final int optimizer;
+
+ /** Learned or random Stiefel transform */
+ public final StiefelTransform stiefelTransform;
+
+ /** Debug hook to disable learning paths */
+ @VisibleForTesting
+ public boolean learn = true;
+
+ // ---------------------------------------------------------------------
+ // Quantizer instances (cached)
+ // ---------------------------------------------------------------------
+
+ private final BinaryQuantizer randomQuantizer = new RandomBinaryQuantizer();
+ private final BinaryQuantizer itqQuantizer = new ItqBinaryQuantizer();
+ private final BinaryQuantizer landingQuantizer = new LandingBinaryQuantizer();
+
+ private BinaryQuantizer quantizer() {
+ switch (optimizer) {
+ case RANDOM:
+ return randomQuantizer;
+ case ITQ:
+ return itqQuantizer;
+ case LANDING:
+ return landingQuantizer;
+ default:
+ throw new IllegalStateException("Unknown optimizer " + optimizer);
+ }
+ }
+
+ /**
+ * Class constructor.
+ * @param originalDim the dimensionality of the vectors to be quantized
+ * @param globalMean the mean of the database (its average vector)
+ */
+ private AsymmetricHashing(int originalDim,
+ int encodedBits,
+ int quantizedDim,
+ int optimizer,
+ VectorFloat> globalMean,
+ StiefelTransform stiefelTransform) {
+ this.originalDimension = originalDim;
+ this.encodedBits = encodedBits;
+ this.quantizedDim = quantizedDim;
+ this.optimizer = optimizer;
+ this.stiefelTransform = Objects.requireNonNull(stiefelTransform, "stiefelTransform");
+ this.globalMean = Objects.requireNonNull(globalMean, "globalMean");
+
+ if (globalMean.length() != originalDimension) {
+ throw new IllegalArgumentException(
+ "Global mean length does not match original dimension");
+ }
+ if (quantizedDim < 1) {
+ throw new IllegalArgumentException("quantizedDim must be > 0");
+ }
+ }
+
+ /**
+ * Initialize ASH index-wide parameters.
+ *
+ * @param ravv the vectors to quantize
+ * @param optimizer the optimizer to use
+ * @param encodedBits the number of bits used to encode vector, including the header
+ */
+ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
+ int optimizer,
+ int encodedBits) {
+
+ final int quantizedDim = encodedBits - HEADER_BITS;
+ if (quantizedDim < 1) {
+ throw new IllegalArgumentException(
+ "Illegal ASH quantized dimensionality: " + quantizedDim);
+ }
+
+ var ravvCopy = ravv.threadLocalSupplier().get();
+ int originalDim = ravvCopy.getVector(0).length();
+
+ // Compute global mean
+ var globalMean = vectorTypeSupport.createFloatVector(originalDim);
+ for (int i = 0; i < ravvCopy.size(); i++) {
+ VectorUtil.addInPlace(globalMean, ravvCopy.getVector(i));
+ }
+ VectorUtil.scale(globalMean, 1.0f / ravvCopy.size());
+
+ StiefelTransform stiefelTransform;
+ if (optimizer == RANDOM) {
+ stiefelTransform = runWithoutTraining(originalDim, quantizedDim, new Random(42));
+ } else if (optimizer == ITQ) {
+ throw new IllegalArgumentException("ITQ optimizer not implemented yet");
+ } else if (optimizer == LANDING) {
+ throw new IllegalArgumentException("LANDING optimizer not implemented yet");
+ } else {
+ throw new IllegalArgumentException("Unknown optimizer " + optimizer);
+ }
+
+ return new AsymmetricHashing(
+ originalDim,
+ encodedBits,
+ quantizedDim,
+ optimizer,
+ globalMean,
+ stiefelTransform
+ );
+ }
+
+ /**
+ * Loads an ASH instance from a RandomAccessReader.
+ * This must mirror the ASH header writer exactly.
+ */
+ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
+
+ int magic = in.readInt();
+ if (magic != MAGIC) {
+ throw new IllegalArgumentException(
+ String.format("Invalid ASH magic: 0x%08X", magic)
+ );
+ }
+
+ int version = in.readInt();
+ if (version > OnDiskGraphIndex.CURRENT_VERSION) {
+ throw new IllegalArgumentException("Unsupported ASH version " + version);
+ }
+
+ int originalDimension = in.readInt();
+ int encodedBits = in.readInt();
+ int quantizedDim = in.readInt();
+ int optimizer = in.readInt();
+
+ int globalMeanLength = in.readInt();
+ if (globalMeanLength <= 0) {
+ throw new IOException("ASH header missing globalMean");
+ }
+ VectorFloat> globalMean = vectorTypeSupport.readFloatVector(in, globalMeanLength);
+
+ // NOTE: StiefelTransform loading will be added in Step 2.
+ // For now, construct a placeholder or re-run the RANDOM initializer.
+ StiefelTransform stiefelTransform;
+
+ if (optimizer == RANDOM) {
+ // Deterministic re-construction for now
+ stiefelTransform = runWithoutTraining(
+ originalDimension,
+ quantizedDim,
+ new Random(42)
+ );
+ } else {
+ throw new IllegalArgumentException(
+ "ASH optimizer " + optimizer + " requires learned transform; not yet supported in load()"
+ );
+ }
+
+ return new AsymmetricHashing(
+ originalDimension,
+ encodedBits,
+ quantizedDim,
+ optimizer,
+ globalMean,
+ stiefelTransform
+ );
+ }
+
+
+ /**
+ * TODO confirm handling versioning correctly
+ * TODO confirm handling magic number correctly
+ * Writes the ASH index header to DataOutput (used for on-disk indexing)
+ * @param out the DataOutput into which to write the object
+ * @throws IOException if there is a problem writing to out.
+ */
+ public void write(DataOutput out, int version) throws IOException {
+ if (version > OnDiskGraphIndex.CURRENT_VERSION) {
+ throw new IllegalArgumentException("Unsupported version " + version);
+ }
+
+ out.writeInt(MAGIC);
+ out.writeInt(version);
+
+ out.writeInt(originalDimension);
+ out.writeInt(encodedBits);
+ out.writeInt(quantizedDim);
+ out.writeInt(optimizer);
+
+ out.writeInt(globalMean.length());
+ vectorTypeSupport.writeFloatVector(out, globalMean);
+
+ // NOTE: stiefelTransform serialization will be added later
+ }
+
+ // ---------------------------------------------------------------------
+ // Index-level quantization entry point
+ // ---------------------------------------------------------------------
+
+ public void quantizeVector(VectorFloat> vector,
+ float residualNorm,
+ float dotWithLandmark,
+ short landmark,
+ QuantizedVector dest) {
+
+ if (quantizedDim > originalDimension) {
+ throw new IllegalArgumentException(
+ "quantizedDim (" + quantizedDim +
+ ") exceeds original dimension (" + originalDimension + ")");
+ }
+
+ QuantizedVector.quantizeTo(
+ vector,
+ residualNorm,
+ dotWithLandmark,
+ landmark,
+ quantizedDim,
+ globalMean,
+ stiefelTransform,
+ quantizer(),
+ dest
+ );
+ }
+
+ // ---------------------------------------------------------------------
+ // Binary quantization strategies
+ // ---------------------------------------------------------------------
+
+ public interface BinaryQuantizer {
+ void quantizeBody(VectorFloat> vector,
+ VectorFloat> globalMean,
+ int quantizedDim,
+ StiefelTransform stiefel,
+ long[] outWords);
+ }
+
+ public static final class RandomBinaryQuantizer implements BinaryQuantizer {
+
+ @Override
+ public void quantizeBody(VectorFloat> vector,
+ VectorFloat> globalMean,
+ int quantizedDim,
+ StiefelTransform stiefel,
+ long[] outWords) {
+
+ RealMatrix A = stiefel.A; // quantizedDim × originalDim
+ int originalDim = A.getColumnDimension();
+
+ final int words = QuantizedVector.wordsForDims(quantizedDim);
+
+ for (int w = 0; w < words; w++) {
+ long bits = 0L;
+ int base = w << 6;
+ int rem = Math.min(64, quantizedDim - base);
+
+ for (int j = 0; j < rem; j++) {
+ int bitIndex = base + j;
+
+ // Compute y_j = dot(A[j], x - μ)
+ double acc = 0.0;
+ for (int d = 0; d < originalDim; d++) {
+ acc += A.getEntry(bitIndex, d)
+ * (vector.get(d) - globalMean.get(d));
+ }
+
+ if (acc > 0.0) {
+ bits |= (1L << j);
+ }
+ }
+
+ outWords[w] = bits;
+ }
+ }
+ }
+
+ public static final class ItqBinaryQuantizer implements BinaryQuantizer {
+ @Override
+ public void quantizeBody(VectorFloat> vector,
+ VectorFloat> globalMean,
+ int quantizedDim,
+ StiefelTransform stiefel,
+ long[] outWords) {
+ throw new UnsupportedOperationException("ITQ not implemented");
+ }
+ }
+
+ public static final class LandingBinaryQuantizer implements BinaryQuantizer {
+ @Override
+ public void quantizeBody(VectorFloat> vector,
+ VectorFloat> globalMean,
+ int quantizedDim,
+ StiefelTransform stiefel,
+ long[] outWords) {
+ throw new UnsupportedOperationException("LANDING not implemented");
+ }
+ }
+
+ // ---------------------------------------------------------------------
+ // Quantized vector (per-vector data)
+ // ---------------------------------------------------------------------
+
+ public static class QuantizedVector {
+ public float residualNorm;
+ public float dotWithLandmark;
+ public short landmark;
+ public long[] binaryVector;
+
+ private QuantizedVector(float residualNorm,
+ float dotWithLandmark,
+ short landmark,
+ long[] binaryVector) {
+ this.residualNorm = residualNorm;
+ this.dotWithLandmark = dotWithLandmark;
+ this.landmark = landmark;
+ this.binaryVector = binaryVector;
+ }
+
+ public static int wordsForDims(int quantizedDim) {
+ return (quantizedDim + 63) >>> 6;
+ }
+
+ public static QuantizedVector createEmpty(int quantizedDim) {
+ return new QuantizedVector(
+ 0f, 0f, (short) 0,
+ new long[wordsForDims(quantizedDim)]);
+ }
+
+ public static int serializedSizeBytes(int quantizedDim) {
+ return Float.BYTES + Float.BYTES + Short.BYTES
+ + wordsForDims(quantizedDim) * Long.BYTES;
+ }
+
+ static void quantizeTo(VectorFloat> vector,
+ float residualNorm,
+ float dotWithLandmark,
+ short landmark,
+ int quantizedDim,
+ VectorFloat> globalMean,
+ StiefelTransform stiefel,
+ BinaryQuantizer quantizer,
+ QuantizedVector dest) {
+
+ dest.residualNorm = residualNorm;
+ dest.dotWithLandmark = dotWithLandmark;
+ dest.landmark = landmark;
+
+ int words = wordsForDims(quantizedDim);
+ if (dest.binaryVector.length < words) {
+ throw new IllegalArgumentException("binaryVector too short");
+ }
+
+ Arrays.fill(dest.binaryVector, 0, words, 0L);
+ quantizer.quantizeBody(vector, globalMean, quantizedDim, stiefel, dest.binaryVector);
+ }
+
+ public void write(DataOutput out, int quantizedDim) throws IOException {
+ out.writeFloat(residualNorm);
+ out.writeFloat(dotWithLandmark);
+ out.writeInt(landmark);
+
+ int words = wordsForDims(quantizedDim);
+ for (int i = 0; i < words; i++) {
+ out.writeLong(binaryVector[i]);
+ }
+ }
+
+ public static QuantizedVector load(RandomAccessReader in,
+ int quantizedDim) throws IOException {
+ float residualNorm = in.readFloat();
+ float dotWithLandmark = in.readFloat();
+ short landmark = (short) in.readInt();
+
+ int words = wordsForDims(quantizedDim);
+ long[] body = new long[words];
+ for (int i = 0; i < words; i++) {
+ body[i] = in.readLong();
+ }
+
+ return new QuantizedVector(residualNorm, dotWithLandmark, landmark, body);
+ }
+
+ public static void loadInto(RandomAccessReader in,
+ QuantizedVector dest,
+ int quantizedDim) throws IOException {
+ dest.residualNorm = in.readFloat();
+ dest.dotWithLandmark = in.readFloat();
+ dest.landmark = (short) in.readInt();
+
+ int words = wordsForDims(quantizedDim);
+ for (int i = 0; i < words; i++) {
+ dest.binaryVector[i] = in.readLong();
+ }
+ }
+
+ /**
+ * Compares two ASH quantized vectors for exact encoded equality.
+ *
+ *
+ * Equality is defined as bitwise equality of all encoded components:
+ *
+ * - Binary code payload
+ * - Associated landmark index
+ * - Auxiliary scalar values (residual norm, dot-with-landmark)
+ *
+ *
+ *
+ * Floating-point values are compared using their raw bit representations
+ * rather than tolerance-based comparisons. This is intentional: quantized
+ * vectors are expected to be deterministic outputs of the encoder, and
+ * approximate equality would be ambiguous and unsafe for hashing.
+ *
+ */
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+
+ QuantizedVector that = (QuantizedVector) o;
+ return Float.floatToIntBits(residualNorm) == Float.floatToIntBits(that.residualNorm)
+ && Float.floatToIntBits(dotWithLandmark) == Float.floatToIntBits(that.dotWithLandmark)
+ && landmark == that.landmark
+ && Arrays.equals(binaryVector, that.binaryVector);
+ }
+
+ /**
+ * Hash code consistent with {@link #equals(Object)}.
+ *
+ *
+ * This hash reflects the full encoded state of the vector and is suitable
+ * for use in hash-based collections during testing, debugging, or caching.
+ *
+ */
+ @Override
+ public int hashCode() {
+ int result = Integer.hashCode(Float.floatToIntBits(residualNorm));
+ result = 31 * result + Integer.hashCode(Float.floatToIntBits(dotWithLandmark));
+ result = 31 * result + Short.hashCode(landmark);
+ result = 31 * result + Arrays.hashCode(binaryVector);
+ return result;
+ }
+
+ }
+
+ /**
+ * Private helper to provide byte storage for binary body
+ * @param nBits the number of bits to store
+ * @return the byte sequence
+ */
+ private static ByteSequence> createBodyByteSequenceForBits(int nBits) {
+ int nBytes = (nBits + 7) >>> 3;
+ return vectorTypeSupport.createByteSequence(nBytes);
+ }
+
+ /**
+ * Orthonormal linear transform on the Stiefel manifold.
+ *
+ * W : originalDim × quantizedDim
+ * A : quantizedDim × originalDim, where A = Wᵀ
+ *
+ * The transform is applied as:
+ * y = A · (x − μ)
+ */
+ public static final class StiefelTransform {
+
+ /** Forward projection (column-orthonormal). */
+ public final RealMatrix W;
+
+ /** Transpose of W, used for projection. */
+ public final RealMatrix A;
+
+ public StiefelTransform(RealMatrix W) {
+ this.W = W;
+ this.A = W.transpose();
+ }
+ }
+
+ /**
+ * RANDOM (untrained) Stiefel initialization.
+ *
+ * Generates a random orthonormal projection from R^originalDim → R^quantizedDim
+ * by sampling a Gaussian matrix and orthonormalizing via QR decomposition.
+ *
+ * A deterministic sign normalization is applied so results are reproducible
+ * across runs and linear algebra backends.
+ *
+ * @param originalDim dimensionality of the input vectors
+ * @param quantizedDim dimensionality of the binary embedding
+ * @param rng random number generator
+ */
+ public static StiefelTransform runWithoutTraining(int originalDim,
+ int quantizedDim,
+ Random rng) {
+
+ if (quantizedDim <= 0 || quantizedDim > originalDim) {
+ throw new IllegalArgumentException(
+ "Invalid quantizedDim=" + quantizedDim +
+ " for originalDim=" + originalDim
+ );
+ }
+
+ // Sample Gaussian matrix G ∈ R^{originalDim × quantizedDim}
+ double[][] gaussian = new double[originalDim][quantizedDim];
+ for (int i = 0; i < originalDim; i++) {
+ for (int j = 0; j < quantizedDim; j++) {
+ gaussian[i][j] = rng.nextGaussian();
+ }
+ }
+
+ RealMatrix G = MatrixUtils.createRealMatrix(gaussian);
+
+ // Thin QR decomposition: G = Q R
+ QRDecomposition qr = new QRDecomposition(G);
+ RealMatrix Q = qr.getQ(); // originalDim × quantizedDim
+ RealMatrix R = qr.getR();
+
+ // Normalize QR sign ambiguity for determinism (RANDOM only)
+ int diag = Math.min(R.getRowDimension(), R.getColumnDimension());
+ for (int j = 0; j < diag; j++) {
+ if (R.getEntry(j, j) < 0.0) {
+ for (int i = 0; i < Q.getRowDimension(); i++) {
+ Q.setEntry(i, j, -Q.getEntry(i, j));
+ }
+ }
+ }
+
+ return new StiefelTransform(Q);
+ }
+
+ @Override
+ @Deprecated
+ public CompressedVectors createCompressedVectors(Object[] compressedVectors) {
+ throw new UnsupportedOperationException("Deprecated createCompressedVectors() method not supported");
+ }
+
+
+ /**
+ * Encodes all vectors using ASH, honoring the caller-provided ForkJoinPool.
+ *
+ * Concurrency model:
+ *
+ * - The work is partitioned into a small number of contiguous ordinal ranges ("chunks").
+ * - Each chunk is processed sequentially by a single task.
+ * - Chunks are executed in parallel using the provided {@code simdExecutor}.
+ *
+ *
+ *
+ * This design intentionally avoids {@code IntStream.parallel()} and other
+ * constructs that implicitly use {@link ForkJoinPool#commonPool()}.
+ * All parallelism is derived exclusively from the supplied executor,
+ * ensuring predictable CPU usage and isolation from unrelated workloads.
+ *
+ *
+ * Correctness guarantees:
+ *
+ * - Each vector ordinal {@code i} is written exactly once to {@code out[i]}.
+ * - The final output array is ordered by ordinal, independent of execution order.
+ * - Missing vectors ({@code null}) are encoded as zero vectors, per the
+ * {@link VectorCompressor} contract.
+ * - Safe publication is guaranteed by joining all submitted tasks
+ * before returning.
+ *
+ *
+ *
+ * This implementation favors clarity and correctness over maximal parallel
+ * granularity. More aggressive chunking or SIMD-aware partitioning can be
+ * layered on later without changing observable semantics.
+ *
+ */
+
+ @Override
+ public ASHVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) {
+ final int n = ravv.size();
+
+ // The caller controls parallelism; we derive a bounded number of tasks
+ // directly from the executor rather than using parallel streams.
+ final int parallelism = Math.max(1, simdExecutor.getParallelism());
+
+ // Each task processes a contiguous range of ordinals.
+ // Chunking by contiguous ranges preserves cache locality and
+ // guarantees each ordinal is written exactly once.
+ final int chunkSize = (n + parallelism - 1) / parallelism;
+
+ final QuantizedVector[] out = new QuantizedVector[n];
+ final var ravvSupplier = ravv.threadLocalSupplier();
+
+ // NOTE:
+ // We intentionally do NOT use IntStream.parallel() here.
+ // Parallel streams always execute on ForkJoinPool.commonPool(),
+ // which would ignore the caller-provided executor and risk
+ // oversubscribing CPU resources.
+ var tasks = new java.util.ArrayList>(parallelism);
+
+ for (int t = 0; t < parallelism; t++) {
+ final int start = t * chunkSize;
+ final int end = Math.min(n, start + chunkSize);
+ if (start >= end) {
+ break;
+ }
+
+ tasks.add(simdExecutor.submit(() -> {
+ // Each task uses its own thread-local RAVV view.
+ var localRavv = ravvSupplier.get();
+
+ for (int i = start; i < end; i++) {
+ VectorFloat> v = localRavv.getVector(i);
+
+ // Contract: missing vectors are encoded as zero vectors.
+ if (v != null) {
+ out[i] = encode(v);
+ } else {
+ out[i] = QuantizedVector.createEmpty(quantizedDim);
+ }
+ }
+ }));
+ }
+
+ // Join establishes a happens-before relationship, ensuring safe publication
+ // of all writes to the output array.
+ for (var task : tasks) {
+ task.join();
+ }
+
+ return new ASHVectors(this, out);
+ }
+
+
+ /**
+ * Encodes the input vector using ASH.
+ * @return the header and binary vector payload
+ */
+ @Override
+ public QuantizedVector encode(VectorFloat> vector) {
+ var qv = QuantizedVector.createEmpty(this.quantizedDim);
+ encodeTo(vector, qv);
+ return qv;
+ }
+
+ // NOTE:
+ // This is the single-landmark baseline configuration.
+ // All vectors are encoded relative to the global mean (landmark 0).
+ // The design intentionally leaves room for multiple landmarks in the future
+ // without changing the on-disk format.
+ @Override
+ public void encodeTo(VectorFloat> vector, QuantizedVector dest) {
+ // Single-landmark baseline: global mean
+ // Landmark index is always 0 in this mode
+ short landmark = 0;
+
+ // Compute dot product with the landmark
+ float dotWithLandmark = VectorUtil.dotProduct(vector, globalMean);
+
+ // Compute residual norm ||x - μ||₂
+ // NOTE:
+ // We store the true L2 norm (not squared) to match ASH semantics and the paper.
+ // squareL2Distance is used for efficiency; the sqrt cost is negligible relative
+ // to projection and hashing.
+ float sqDist = VectorUtil.squareL2Distance(vector, globalMean);
+ float residualNorm = (float) Math.sqrt(sqDist);
+
+ // Binary encoding uses sign(A · (x − μ)) internally
+ quantizeVector(vector, residualNorm, dotWithLandmark, landmark, dest);
+ }
+
+ @Override
+ public int compressorSize() {
+ // Must exactly match write(out, version)
+ int size = 0;
+ size += Integer.BYTES; // MAGIC
+ size += Integer.BYTES; // version
+ size += Integer.BYTES; // originalDimension
+ size += Integer.BYTES; // encodedBits
+ size += Integer.BYTES; // quantizedDim
+ size += Integer.BYTES; // optimizer
+ size += Integer.BYTES; // globalMean length
+ size += Float.BYTES * globalMean.length();
+ return size;
+ }
+
+ @Override
+ public int compressedVectorSize() {
+ int words = QuantizedVector.wordsForDims(quantizedDim);
+ return Float.BYTES // residualNorm
+ + Float.BYTES // dotWithLandmark
+ + Integer.BYTES // landmark
+ + words * Long.BYTES; // binary payload
+ }
+
+ @Override
+ public double reconstructionError(VectorFloat> vector) {
+ throw new UnsupportedOperationException("Reconstruction error not defined for ASH");
+ }
+
+ @Override
+ public long ramBytesUsed() {
+ long bytes = 0L;
+
+ if (globalMean != null) {
+ bytes += globalMean.ramBytesUsed();
+ }
+
+ if (stiefelTransform != null) {
+ bytes += estimateMatrixBytes(stiefelTransform.W);
+ bytes += estimateMatrixBytes(stiefelTransform.A);
+ }
+
+ return bytes;
+ }
+
+ private static long estimateMatrixBytes(RealMatrix m) {
+ if (m == null) return 0L;
+ int rows = m.getRowDimension();
+ int cols = m.getColumnDimension();
+ return (long) rows * cols * Double.BYTES
+ + (long) rows * Long.BYTES;
+ }
+
+ // Exact, element-wise comparison of VectorFloat contents.
+ // We use Float.floatToIntBits to handle -0.0f, NaN, and ensure
+ // stable, bitwise equality suitable for hashing.
+ private static boolean vectorEquals(VectorFloat> a, VectorFloat> b) {
+ if (a == b) return true;
+ if (a == null || b == null) return false;
+ if (a.length() != b.length()) return false;
+
+ for (int i = 0; i < a.length(); i++) {
+ if (Float.floatToIntBits(a.get(i)) != Float.floatToIntBits(b.get(i))) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ // Hash based on the raw bit representation of each float.
+ // This mirrors vectorEquals and guarantees hash consistency.
+ private static int vectorHash(VectorFloat> v) {
+ if (v == null) return 0;
+ int h = 1;
+ for (int i = 0; i < v.length(); i++) {
+ h = 31 * h + Integer.hashCode(Float.floatToIntBits(v.get(i)));
+ }
+ return h;
+ }
+
+ // Exact comparison of two RealMatrix instances by dimensions and entries.
+ // We intentionally avoid RealMatrix.equals() and tolerance-based checks.
+ // ASH initialization is deterministic, so bitwise equality is expected
+ // and required for robust hashing semantics.
+ private static boolean matrixEquals(RealMatrix a, RealMatrix b) {
+ if (a == b) return true;
+ if (a == null || b == null) return false;
+ if (a.getRowDimension() != b.getRowDimension()
+ || a.getColumnDimension() != b.getColumnDimension()) {
+ return false;
+ }
+
+ for (int i = 0; i < a.getRowDimension(); i++) {
+ for (int j = 0; j < a.getColumnDimension(); j++) {
+ if (Double.doubleToLongBits(a.getEntry(i, j))
+ != Double.doubleToLongBits(b.getEntry(i, j))) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+
+ // Hash code derived from matrix dimensions and raw double bits.
+ // Only the forward projection matrix (W) is hashed; the transpose (A)
+ // is derived and does not need to be hashed separately.
+ private static int matrixHash(RealMatrix m) {
+ if (m == null) return 0;
+ int h = 1;
+ h = 31 * h + m.getRowDimension();
+ h = 31 * h + m.getColumnDimension();
+
+ for (int i = 0; i < m.getRowDimension(); i++) {
+ for (int j = 0; j < m.getColumnDimension(); j++) {
+ long bits = Double.doubleToLongBits(m.getEntry(i, j));
+ h = 31 * h + (int) (bits ^ (bits >>> 32));
+ }
+ }
+ return h;
+ }
+
+ /**
+ * Compares this ASH compressor to another for encoding equivalence.
+ *
+ *
+ * Two compressors are equal if they would produce identical encoded outputs
+ * (including binary codes and auxiliary scalars) for the same input vectors.
+ *
+ *
+ *
+ * Equality therefore requires:
+ *
+ * - Matching dimensional parameters
+ * - Identical optimizer selection
+ * - Bitwise-equal global mean vectors
+ * - Bitwise-equal Stiefel transform matrices
+ *
+ *
+ *
+ * Note that this is a strong notion of equality. We deliberately
+ * avoid approximate floating-point comparisons because ASH initialization
+ * and training are deterministic, and weaker equality would break hashing
+ * and cache correctness.
+ *
+ */
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+
+ AsymmetricHashing that = (AsymmetricHashing) o;
+ return originalDimension == that.originalDimension
+ && encodedBits == that.encodedBits
+ && quantizedDim == that.quantizedDim
+ && optimizer == that.optimizer
+ && vectorEquals(this.globalMean, that.globalMean)
+ && matrixEquals(
+ this.stiefelTransform == null ? null : this.stiefelTransform.W,
+ that.stiefelTransform == null ? null : that.stiefelTransform.W
+ );
+ }
+
+ /**
+ * Hash code consistent with {@link #equals(Object)}.
+ *
+ *
+ * The hash incorporates all parameters and learned state that affect encoding,
+ * including the contents of the global mean and Stiefel transform. This ensures
+ * safe use of ASH compressors as map keys or cache entries.
+ *
+ */
+ @Override
+ public int hashCode() {
+ int result = Integer.hashCode(originalDimension);
+ result = 31 * result + Integer.hashCode(encodedBits);
+ result = 31 * result + Integer.hashCode(quantizedDim);
+ result = 31 * result + Integer.hashCode(optimizer);
+ result = 31 * result + vectorHash(globalMean);
+ result = 31 * result + matrixHash(
+ stiefelTransform == null ? null : stiefelTransform.W
+ );
+ return result;
+ }
+
+ @Override
+ public String toString() {
+ String optimizerName;
+ switch (optimizer) {
+ case RANDOM:
+ optimizerName = "RANDOM";
+ break;
+ case ITQ:
+ optimizerName = "ITQ";
+ break;
+ case LANDING:
+ optimizerName = "LANDING";
+ break;
+ default:
+ optimizerName = "UNKNOWN(" + optimizer + ")";
+ }
+
+ return String.format(
+ "AsymmetricHashing[origDim=%d, encodedBits=%d, quantizedDim=%d, optimizer=%s, landmarks=1]",
+ originalDimension,
+ encodedBits,
+ quantizedDim,
+ optimizerName
+ );
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
new file mode 100644
index 000000000..b9e5ee516
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -0,0 +1,153 @@
+/*
+ * 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;
+
+import io.github.jbellis.jvector.example.util.SiftLoader;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.io.IOException;
+import java.util.List;
+
+import static java.lang.Math.abs;
+
+// this class uses explicit typing instead of `var` for easier reading when excerpted for instructional use
+public class DistancesASH {
+ public static void testASHEncodings(String filenameBase, String filenameQueries, VectorSimilarityFunction vsf) throws IOException {
+ List> vectors = SiftLoader.readFvecs(filenameBase);
+ List> queries = SiftLoader.readFvecs(filenameQueries);
+
+ int dimension = vectors.get(0).length();
+ int encodedDim = dimension; // change later
+ int nQueries = 100;
+ int nVectors = 10_000;
+
+ vectors = vectors.subList(0, nVectors);
+
+ System.out.format("\t%d base and %d query vectors loaded, dimensions %d%n",
+ vectors.size(), queries.size(), vectors.get(0).length());
+
+ // Generate a BQ for random vectors
+ var ravv = new ListRandomAccessVectorValues(vectors, dimension);
+ var ash = AsymmetricHashing.initialize(ravv, AsymmetricHashing.RANDOM, encodedDim);
+
+ // Compress the vectors
+ long startTime = System.nanoTime();
+ var ashVecs = ash.encodeAll(ravv);
+ long endTime = System.nanoTime();
+ double duration = (double) (endTime - startTime) / 1_000_000_000;
+ System.out.println("\tEncoding took " + duration + " seconds");
+
+ // compare the encoded similarities to the raw
+ double distanceError = 0;
+ for (int i = 0; i < nQueries; i++) {
+ var q = queries.get(i);
+ if (VectorUtil.dotProduct(q, q) == 0) {
+ continue;
+ }
+ var f = ashVecs.scoreFunctionFor(q, vsf);
+
+ for (int j = 0; j < nVectors; j++) {
+ var v = vectors.get(j);
+ distanceError += abs(f.similarityTo(j) - vsf.compare(q, v));
+ }
+ }
+ distanceError /= nQueries * nVectors;
+
+ System.out.println("\t" + vsf + " error " + distanceError);
+
+
+ float dummyAccumulator = 0;
+
+ startTime = System.nanoTime();
+ for (int i = 0; i < nQueries; i++) {
+ var q = queries.get(i);
+ if (VectorUtil.dotProduct(q, q) == 0) {
+ continue;
+ }
+ var f = ashVecs.scoreFunctionFor(q, vsf);
+
+ for (int j = 0; j < nVectors; j++) {
+ dummyAccumulator += f.similarityTo(j);
+ }
+ }
+ endTime = System.nanoTime();
+ duration = (double) (endTime - startTime) / 1_000_000_000;
+ System.out.println("\tASH Distance computations took " + duration + " seconds");
+
+ startTime = System.nanoTime();
+ for (int i = 0; i < nQueries; i++) {
+ var q = queries.get(i);
+ if (VectorUtil.dotProduct(q, q) == 0) {
+ continue;
+ }
+
+ for (int j = 0; j < nVectors; j++) {
+ var v = vectors.get(j);
+ dummyAccumulator += vsf.compare(q, v);
+ }
+ }
+ endTime = System.nanoTime();
+ duration = (double) (endTime - startTime) / 1_000_000_000;
+ System.out.println("\tFloat Distance computations took " + duration + " seconds");
+
+ System.out.println("\tdummyAccumulator: " + dummyAccumulator);
+ System.out.println("--");
+ }
+
+ public static void runSIFT() throws IOException {
+ System.out.println("Running siftsmall");
+
+ var baseVectors = "siftsmall/siftsmall_base.fvecs";
+ var queryVectors = "siftsmall/siftsmall_query.fvecs";
+ testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ }
+
+ public static void runADA() throws IOException {
+ System.out.println("Running ada_002");
+
+ var baseVectors = "./fvec/wikipedia_squad/100k/ada_002_100000_base_vectors.fvec";
+ var queryVectors = "./fvec/wikipedia_squad/100k/ada_002_100000_query_vectors_10000.fvec";
+ testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ }
+
+ public static void runColbert() throws IOException {
+ System.out.println("Running colbertv2");
+
+ var baseVectors = "./fvec/wikipedia_squad/1M/colbertv2.0_128_base_vectors_1000000.fvec";
+ var queryVectors = "./fvec/wikipedia_squad/1M/colbertv2.0_128_query_vectors_100000.fvec";
+ testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ }
+
+ public static void runOpenai3072() throws IOException {
+ System.out.println("Running text-embedding-3-large_3072");
+
+ var baseVectors = "./fvec/wikipedia_squad/100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
+ var queryVectors = "./fvec/wikipedia_squad/100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
+ testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ }
+
+ public static void main(String[] args) throws IOException {
+ runSIFT();
+ runADA();
+ runColbert();
+// runOpenai3072();
+ }
+}
\ No newline at end of file
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
new file mode 100644
index 000000000..5ada6365b
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
@@ -0,0 +1,43 @@
+package io.github.jbellis.jvector.example;
+
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
+import org.apache.commons.math3.linear.*;
+import java.util.Random;
+
+public class TestAshRunWithoutTraining {
+
+ public static void main(String[] args) {
+ int d = 1024; // input dimension
+ int nDims = 128; // target dimension
+
+ // Dummy data matrix X only to scrape shape
+ RealMatrix X = MatrixUtils.createRealMatrix(10, d); // (10 × 1024)
+
+ Random rng = new Random(42);
+
+ // Method under test
+ AsymmetricHashing.StiefelTransform stiefelTransform =
+ AsymmetricHashing.runWithoutTraining(X, nDims, rng);
+
+ // Inspect shapes and a few entries
+ RealMatrix A = stiefelTransform.A;
+ RealMatrix W = stiefelTransform.W;
+
+ System.out.println("A shape = " + A.getRowDimension() + " x " + A.getColumnDimension());
+ System.out.println("W shape = " + W.getRowDimension() + " x " + W.getColumnDimension());
+
+ // Sanity check....
+ System.out.println("First few entries of A:");
+ for (int i = 0; i < Math.min(3, A.getRowDimension()); i++) {
+ for (int j = 0; j < Math.min(3, A.getColumnDimension()); j++) {
+ System.out.printf("%.4f ", A.getEntry(i, j));
+ }
+ System.out.println();
+ }
+
+ // Check orthogonality: WᵀW ≈ I
+ RealMatrix I = W.transpose().multiply(W);
+ System.out.printf("I(0,0)=%.5f, I(0,1)=%.5f%n",
+ I.getEntry(0, 0), I.getEntry(0, 1));
+ }
+}
From 7a604227a0e08a6da76e6b1fe6d4814ca974c05a Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Wed, 24 Dec 2025 12:51:40 -0800
Subject: [PATCH 002/243] First working version with scorer. Other minor
improvements.
---
.../jvector/quantization/ASHScorer.java | 135 ++++++++++
.../jvector/quantization/ASHVectors.java | 51 +---
.../quantization/AsymmetricHashing.java | 125 ++++++---
.../jbellis/jvector/example/DistancesASH.java | 249 +++++++++++++-----
.../example/TestAshRunWithoutTraining.java | 9 +-
5 files changed, 423 insertions(+), 146 deletions(-)
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
new file mode 100644
index 000000000..93f35182f
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
@@ -0,0 +1,135 @@
+/*
+ * 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.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+/**
+ * Query-time scorer for ASH vectors.
+ * Single-landmark case (C = 1).
+ * NOTE: Only DOT_PRODUCT is supported for now.
+ */
+public final class ASHScorer {
+ private final AsymmetricHashing ash;
+
+ public ASHScorer(AsymmetricHashing ash) {
+ this.ash = ash;
+ assert ash.stiefelTransform.rows == ash.quantizedDim
+ : "Stiefel rows must equal quantizedDim";
+ }
+
+ public ASHScoreFunction scoreFunctionFor(VectorFloat> query, VectorSimilarityFunction similarityFunction) {
+ if (similarityFunction != VectorSimilarityFunction.DOT_PRODUCT) {
+ throw new UnsupportedOperationException(
+ "ASH scorer supports DOT_PRODUCT only (requested " + similarityFunction + ")"
+ );
+ }
+ return dotProductScoreFunctionFor(query);
+ }
+
+ /**
+ * Scorer for dot product approximation.
+ */
+ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
+ final int d = ash.quantizedDim;
+
+ // Compute q_centered = q - μ
+ // (avoid allocation by doing subtraction on the fly during projection)
+ final VectorFloat> mu = ash.globalMean;
+
+ final double[][] A = ash.stiefelTransform.AData; // shape: d × originalDim
+ final int originalDim = ash.stiefelTransform.cols;
+ final float[] muArr = new float[originalDim];
+ for (int k = 0; k < originalDim; k++) muArr[k] = mu.get(k);
+
+ final float[] qArr = new float[originalDim];
+ for (int k = 0; k < originalDim; k++) qArr[k] = query.get(k);
+
+ // Project query: tildeQ[j] = (A (q - μ))[j]
+ final float[] tildeQ = new float[d];
+ for (int j = 0; j < d; j++) {
+ final double[] Aj = A[j];
+ double acc = 0.0;
+ for (int k = 0; k < originalDim; k++) {
+ acc += Aj[k] * (qArr[k] - muArr[k]);
+ }
+ tildeQ[j] = (float) acc;
+ }
+
+ // = sum_j tildeQ[j]
+ float tmpSum = 0f;
+ for (int j = 0; j < d; j++) {
+ tmpSum += tildeQ[j];
+ }
+ final float sumTildeQ = tmpSum;
+
+ //
+ final float dotQMu = VectorUtil.dotProduct(query, mu);
+
+ // ||μ||^2
+ final float muNormSq = VectorUtil.dotProduct(mu, mu);
+
+ // Precompute invSqrtD once
+ final float invSqrtD = (float) (1.0 / Math.sqrt(d));
+
+ return (AsymmetricHashing.QuantizedVector v) -> {
+ // Single landmark baseline: ignore v.landmark for now (assumed 0)
+ final float residualNorm = v.residualNorm; // ||x - μ||₂
+ final float dotXMu = v.dotWithLandmark; //
+
+ // scale = d^{-1/2} * ||x - μ||
+ final float scale = invSqrtD * residualNorm;
+
+ // masked-add term: where bin ∈ {0,1}^d
+ // This is sum of tildeQ[j] over set bits in v.binaryVector.
+ float maskedAdd = 0f;
+
+ // v.binaryVector stores bits packed in longs, little-endian per word (bit j corresponds to dimension base+j)
+ final long[] bits = v.binaryVector;
+ int base = 0;
+ for (int w = 0; w < bits.length && base < d; w++, base += 64) {
+ long word = bits[w];
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = base + bit;
+ if (idx < d) maskedAdd += tildeQ[idx];
+ word &= (word - 1);
+ }
+ }
+
+ // ASH dot-product approximation (single landmark, {0,1} encoding):
+ //
+ // scale * (2⟨q̃, b⟩ − ⟨q̃, 1⟩) + ⟨q, μ⟩ + (⟨x, μ⟩ − ||μ||²)
+ //
+ // Notes:
+ // - b ∈ {0,1}^d is the binary code (converted from the paper’s {−1,+1} form).
+ // - scale = ||x − μ|| / √d (pure Eq. 11 baseline; no bias correction).
+ // - ⟨q, μ⟩ is a query-only offset and MUST NOT be scaled.
+ // - Scaling ⟨q, μ⟩ leads to large systematic error (see paper Eq. 11 regrouping).
+ return scale * (2f * maskedAdd - sumTildeQ) + dotQMu + (dotXMu - muNormSq);
+
+ };
+ }
+
+ @FunctionalInterface
+ public interface ASHScoreFunction {
+ float similarityTo(AsymmetricHashing.QuantizedVector vector2);
+ }
+}
+
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index b1f822ea2..dcfe63a39 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -49,17 +49,12 @@
* Define reranking strategy (likely via full vectors)
* Support multi-landmark ASH (this is intentionally deferred)
*
- *
- *
- * NOTE:
- * This class avoids PQ/NVQ feature plumbing and graph integration by design.
- * Those concerns belong to the scoring plane and will be addressed after
- * baseline performance is established.
*/
public class ASHVectors implements CompressedVectors {
final AsymmetricHashing ash;
final AsymmetricHashing.QuantizedVector[] compressedVectors;
+ final ASHScorer scorer;
/**
* Initialize ASHVectors with an array of ASH-compressed vectors.
@@ -71,6 +66,7 @@ public ASHVectors(AsymmetricHashing ash,
AsymmetricHashing.QuantizedVector[] compressedVectors) {
this.ash = ash;
this.compressedVectors = compressedVectors;
+ this.scorer = new ASHScorer(ash);
}
@Override
@@ -118,45 +114,16 @@ public static ASHVectors load(RandomAccessReader in) throws IOException {
return new ASHVectors(ash, compressedVectors);
}
- /**
- * Return an approximate score function for the given query.
- *
- *
- * Baseline implementation:
- *
- * - Encode query once using ASH
- * - Score candidates using pure Hamming distance
- *
- *
- *
- * TODO:
- *
- * - Replace pure Hamming with full ASH scoring
- * - Incorporate landmark scalars (dotWithLandmark, residualNorm)
- * - Respect VectorSimilarityFunction semantics where applicable
- *
- */
@Override
public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(
VectorFloat> query,
VectorSimilarityFunction similarityFunction) {
- // Encode query once
- var qv = AsymmetricHashing.QuantizedVector.createEmpty(ash.quantizedDim);
- ash.encodeTo(query, qv);
- final long[] queryBits = qv.binaryVector;
+ // Scorer follows ASH paper (single landmark, DOT_PRODUCT only)
+ final ASHScorer.ASHScoreFunction f = scorer.scoreFunctionFor(query, similarityFunction);
- return node -> {
- var v = compressedVectors[node];
-
- // TODO:
- // - Replace with proper ASH scoring.
- // - Consider affine scaling or normalization.
- int hamming = VectorUtil.hammingDistance(queryBits, v.binaryVector);
-
- // Higher similarity is better
- return -hamming;
- };
+ // Wrap QuantizedVector scorer into ordinal-based score function
+ return node -> f.similarityTo(compressedVectors[node]);
}
/**
@@ -177,11 +144,7 @@ public ScoreFunction.ApproximateScoreFunction precomputedScoreFunctionFor(
* Diversity-aware scoring is not supported for ASH at this stage.
*
*
- * TODO:
- * - Define whether diversity should be measured via:
- * * Hamming distance
- * * Landmark-aware distance
- * * Full-vector reranking
+ * TODO: Define how diversity should be measured.
*/
@Override
public ScoreFunction.ApproximateScoreFunction diversityFunctionFor(
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 22e1201cf..a69fc9443 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -34,7 +34,10 @@
import java.util.concurrent.ForkJoinPool;
import java.util.Random;
import java.util.stream.IntStream;
-import org.apache.commons.math3.linear.*;
+import org.apache.commons.math3.linear.Array2DRowRealMatrix;
+import org.apache.commons.math3.linear.MatrixUtils;
+import org.apache.commons.math3.linear.QRDecomposition;
+import org.apache.commons.math3.linear.RealMatrix;
import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED;
@@ -47,7 +50,12 @@ public class AsymmetricHashing implements VectorCompressor vector,
public interface BinaryQuantizer {
void quantizeBody(VectorFloat> vector,
VectorFloat> globalMean,
+ float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
long[] outWords);
@@ -303,12 +316,25 @@ public static final class RandomBinaryQuantizer implements BinaryQuantizer {
@Override
public void quantizeBody(VectorFloat> vector,
VectorFloat> globalMean,
+ float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
long[] outWords) {
- RealMatrix A = stiefel.A; // quantizedDim × originalDim
- int originalDim = A.getColumnDimension();
+ final double[][] A = stiefel.AData; // [quantizedDim][originalDim]
+ final int originalDim = A[0].length;
+
+ // Copy vector and mean once (no virtual calls in inner loop)
+ final float[] x = new float[originalDim];
+ final float[] mu = new float[originalDim];
+
+ for (int d = 0; d < originalDim; d++) {
+ x[d] = vector.get(d);
+ mu[d] = globalMean.get(d);
+ }
+
+ // Eq. 6 normalization
+ final double invNorm = (residualNorm > 0f) ? (1.0 / residualNorm) : 0.0;
final int words = QuantizedVector.wordsForDims(quantizedDim);
@@ -320,11 +346,12 @@ public void quantizeBody(VectorFloat> vector,
for (int j = 0; j < rem; j++) {
int bitIndex = base + j;
- // Compute y_j = dot(A[j], x - μ)
+ // y_j = dot(A[j], (x − μ) / ||x − μ||)
double acc = 0.0;
+ final double[] Arow = A[bitIndex];
+
for (int d = 0; d < originalDim; d++) {
- acc += A.getEntry(bitIndex, d)
- * (vector.get(d) - globalMean.get(d));
+ acc += Arow[d] * (x[d] - mu[d]) * invNorm;
}
if (acc > 0.0) {
@@ -341,6 +368,7 @@ public static final class ItqBinaryQuantizer implements BinaryQuantizer {
@Override
public void quantizeBody(VectorFloat> vector,
VectorFloat> globalMean,
+ float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
long[] outWords) {
@@ -352,6 +380,7 @@ public static final class LandingBinaryQuantizer implements BinaryQuantizer {
@Override
public void quantizeBody(VectorFloat> vector,
VectorFloat> globalMean,
+ float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
long[] outWords) {
@@ -364,10 +393,10 @@ public void quantizeBody(VectorFloat> vector,
// ---------------------------------------------------------------------
public static class QuantizedVector {
- public float residualNorm;
- public float dotWithLandmark;
- public short landmark;
- public long[] binaryVector;
+ public float residualNorm; // ||x_i - μ_i*||_2
+ public float dotWithLandmark; //
+ public short landmark; // c_i*
+ public long[] binaryVector; // bin(x̂_i)
private QuantizedVector(float residualNorm,
float dotWithLandmark,
@@ -414,7 +443,14 @@ static void quantizeTo(VectorFloat> vector,
}
Arrays.fill(dest.binaryVector, 0, words, 0L);
- quantizer.quantizeBody(vector, globalMean, quantizedDim, stiefel, dest.binaryVector);
+ quantizer.quantizeBody(
+ vector,
+ globalMean,
+ residualNorm,
+ quantizedDim,
+ stiefel,
+ dest.binaryVector
+ );
}
public void write(DataOutput out, int quantizedDim) throws IOException {
@@ -517,24 +553,46 @@ private static ByteSequence> createBodyByteSequenceForBits(int nBits) {
/**
* Orthonormal linear transform on the Stiefel manifold.
- *
* W : originalDim × quantizedDim
* A : quantizedDim × originalDim, where A = Wᵀ
- *
* The transform is applied as:
* y = A · (x − μ)
*/
public static final class StiefelTransform {
-
/** Forward projection (column-orthonormal). */
public final RealMatrix W;
/** Transpose of W, used for projection. */
public final RealMatrix A;
+ /**
+ * Cached row-major backing for A when available.
+ * This exists purely to avoid RealMatrix.getEntry() in hot loops.
+ * AData must not be mutated after construction.
+ */
+ public final double[][] AData;
+
+ /** Number of rows in A (i.e., quantizedDim = d). */
+ public final int rows;
+
+ /** Number of columns in A (i.e., originalDim). */
+ public final int cols;
+
public StiefelTransform(RealMatrix W) {
this.W = W;
this.A = W.transpose();
+
+ // Fast path: QRDecomposition + MatrixUtils typically produce Array2DRowRealMatrix
+ // Use getDataRef() to avoid copying and avoid getEntry() calls.
+ if (this.A instanceof Array2DRowRealMatrix) {
+ this.AData = ((Array2DRowRealMatrix) this.A).getDataRef();
+ } else {
+ // Fallback: copy once at construction time (still far faster than getEntry per access)
+ this.AData = this.A.getData();
+ }
+
+ this.rows = AData.length;
+ this.cols = AData.length == 0 ? 0 : AData[0].length;
}
}
@@ -699,29 +757,28 @@ public QuantizedVector encode(VectorFloat> vector) {
return qv;
}
- // NOTE:
- // This is the single-landmark baseline configuration.
- // All vectors are encoded relative to the global mean (landmark 0).
- // The design intentionally leaves room for multiple landmarks in the future
- // without changing the on-disk format.
@Override
public void encodeTo(VectorFloat> vector, QuantizedVector dest) {
- // Single-landmark baseline: global mean
- // Landmark index is always 0 in this mode
- short landmark = 0;
+ // NOTE:
+ // Single-landmark baseline configuration (Table 2, C = 1).
+ // All vectors are encoded relative to the global mean μ*.
+ // The landmark id is always 0 in this mode.
- // Compute dot product with the landmark
- float dotWithLandmark = VectorUtil.dotProduct(vector, globalMean);
+ final short landmark = 0;
- // Compute residual norm ||x - μ||₂
- // NOTE:
- // We store the true L2 norm (not squared) to match ASH semantics and the paper.
- // squareL2Distance is used for efficiency; the sqrt cost is negligible relative
- // to projection and hashing.
- float sqDist = VectorUtil.squareL2Distance(vector, globalMean);
- float residualNorm = (float) Math.sqrt(sqDist);
+ // Table 2 header field:
+ final float dotWithLandmark = VectorUtil.dotProduct(vector, globalMean);
+
+ // Table 2 header field: ||x_i − μ_i*||₂
+ // We store the true L2 norm (not squared), matching the paper.
+ final float sqDist = VectorUtil.squareL2Distance(vector, globalMean);
+ final float residualNorm = (float) Math.sqrt(sqDist);
+
+ // Sanity: quantizedDim already accounts for physical header bits (Step 1)
+ // Binary body uses exactly quantizedDim bits.
+ assert quantizedDim > 0;
- // Binary encoding uses sign(A · (x − μ)) internally
+ // Binary body: sign(A · (x − μ))
quantizeVector(vector, residualNorm, dotWithLandmark, landmark, dest);
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index b9e5ee516..695dc21e0 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -18,97 +18,222 @@
import io.github.jbellis.jvector.example.util.SiftLoader;
import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.quantization.CompressedVectors;
+
import io.github.jbellis.jvector.quantization.AsymmetricHashing;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
import io.github.jbellis.jvector.vector.VectorUtil;
import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
import java.io.IOException;
import java.util.List;
+import java.util.stream.Collectors;
+import java.util.ArrayList;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ForkJoinTask;
+
+
import static java.lang.Math.abs;
-// this class uses explicit typing instead of `var` for easier reading when excerpted for instructional use
public class DistancesASH {
- public static void testASHEncodings(String filenameBase, String filenameQueries, VectorSimilarityFunction vsf) throws IOException {
+ public static void testASHEncodings(String filenameBase, String filenameQueries) throws IOException {
+
List> vectors = SiftLoader.readFvecs(filenameBase);
List> queries = SiftLoader.readFvecs(filenameQueries);
+ // ------------------------------------------------------------------
+ // Remove zero-norm vectors (undefined for ASH / Eq. 11)
+ // ------------------------------------------------------------------
+ int beforeBase = vectors.size();
+ List> filteredBase = new ArrayList<>(vectors.size());
+ for (VectorFloat> v : vectors) {
+ if (VectorUtil.dotProduct(v, v) > 0f) {
+ filteredBase.add(v);
+ }
+ }
+ vectors = filteredBase;
+
+ int beforeQuery = queries.size();
+ List> filteredQueries = new ArrayList<>(queries.size());
+ for (VectorFloat> q : queries) {
+ if (VectorUtil.dotProduct(q, q) > 0f) {
+ filteredQueries.add(q);
+ }
+ }
+ queries = filteredQueries;
+
+ System.out.println("\tRemoved " + (beforeBase - vectors.size()) + " zero base vectors and "
+ + (beforeQuery - queries.size()) + " zero query vectors");
+
+ // ASH normalization policy:
+ //
+ // - Base vectors x are NOT globally normalized.
+ // - For encoding, we compute μ on raw x.
+ // - The residual (x − μ) is normalized ONCE inside the binarizer,
+ // producing \hat{x} = (x − μ) / ||x − μ|| (Eq. 6).
+ // - Queries may be L2-normalized in the benchmark (standard practice),
+ // but are NOT normalized inside the encoder.
+ // - No other normalization steps are applied.
+ for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
+
int dimension = vectors.get(0).length();
- int encodedDim = dimension; // change later
- int nQueries = 100;
- int nVectors = 10_000;
+ int encodedBits = dimension; // sweep later if desired
+
+ System.out.println(
+ "\toriginalDim = " + dimension +
+ ", encodedBits = " + encodedBits
+ );
+
+ int nQueries = Math.min(1000, queries.size());
+ int nVectors = Math.min(100_000, vectors.size());
vectors = vectors.subList(0, nVectors);
+ queries = queries.subList(0, nQueries);
+
+ final List> finalQueries = queries;
+ final List> finalVectors = vectors;
- System.out.format("\t%d base and %d query vectors loaded, dimensions %d%n",
- vectors.size(), queries.size(), vectors.get(0).length());
+ System.out.format("\t%d base and %d query vectors loaded, dimension=%d%n",
+ vectors.size(), queries.size(), dimension);
- // Generate a BQ for random vectors
+ // ------------------------------------------------------------------
+ // Build ASH
+ // ------------------------------------------------------------------
var ravv = new ListRandomAccessVectorValues(vectors, dimension);
- var ash = AsymmetricHashing.initialize(ravv, AsymmetricHashing.RANDOM, encodedDim);
+ var ash = AsymmetricHashing.initialize(ravv, AsymmetricHashing.RANDOM, encodedBits);
- // Compress the vectors
long startTime = System.nanoTime();
- var ashVecs = ash.encodeAll(ravv);
+ CompressedVectors ashVecs = ash.encodeAll(ravv);
long endTime = System.nanoTime();
- double duration = (double) (endTime - startTime) / 1_000_000_000;
- System.out.println("\tEncoding took " + duration + " seconds");
-
- // compare the encoded similarities to the raw
- double distanceError = 0;
- for (int i = 0; i < nQueries; i++) {
- var q = queries.get(i);
- if (VectorUtil.dotProduct(q, q) == 0) {
- continue;
- }
- var f = ashVecs.scoreFunctionFor(q, vsf);
+ System.out.println("\tEncoding took " + (endTime - startTime) / 1e9 + " seconds");
- for (int j = 0; j < nVectors; j++) {
- var v = vectors.get(j);
- distanceError += abs(f.similarityTo(j) - vsf.compare(q, v));
- }
+ // ------------------------------------------------------------------
+ // Shared parallelization setup (executor-honoring)
+ // ------------------------------------------------------------------
+ ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool();
+
+ int parallelism = simdExecutor.getParallelism();
+ int chunkSize = Math.max(1, (nQueries + parallelism - 1) / parallelism);
+
+ // ==================================================================
+ // [1] Accuracy run (NOT timed)
+ // ==================================================================
+ List> errorTasks = new ArrayList<>();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ errorTasks.add(simdExecutor.submit(() -> {
+ double localError = 0.0;
+ long localCount = 0;
+
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ ScoreFunction.ApproximateScoreFunction f =
+ ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+
+ for (int j = 0; j < nVectors; j++) {
+ VectorFloat> v = finalVectors.get(j);
+ float trueDot = VectorUtil.dotProduct(q, v);
+ float approxDot = f.similarityTo(j);
+
+ localError += Math.abs(approxDot - trueDot);
+ localCount++;
+ }
+ }
+ return new double[]{localError, localCount};
+ }));
}
- distanceError /= nQueries * nVectors;
- System.out.println("\t" + vsf + " error " + distanceError);
+ double distanceError = 0.0;
+ long count = 0;
+ for (ForkJoinTask t : errorTasks) {
+ double[] r = t.join();
+ distanceError += r[0];
+ count += (long) r[1];
+ }
- float dummyAccumulator = 0;
+ distanceError /= count;
+ System.out.println("\tAverage absolute dot-product error = " + distanceError);
- startTime = System.nanoTime();
- for (int i = 0; i < nQueries; i++) {
- var q = queries.get(i);
- if (VectorUtil.dotProduct(q, q) == 0) {
- continue;
- }
- var f = ashVecs.scoreFunctionFor(q, vsf);
+ // ==================================================================
+ // [2] ASH distance timing (Eq. 11 only)
+ // ==================================================================
+ List> ashTasks = new ArrayList<>();
- for (int j = 0; j < nVectors; j++) {
- dummyAccumulator += f.similarityTo(j);
- }
+ long ashStart = System.nanoTime();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ ashTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
+
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ ScoreFunction.ApproximateScoreFunction f =
+ ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+
+ for (int j = 0; j < nVectors; j++) {
+ localSum += f.similarityTo(j);
+ }
+ }
+ return localSum;
+ }));
}
- endTime = System.nanoTime();
- duration = (double) (endTime - startTime) / 1_000_000_000;
- System.out.println("\tASH Distance computations took " + duration + " seconds");
-
- startTime = System.nanoTime();
- for (int i = 0; i < nQueries; i++) {
- var q = queries.get(i);
- if (VectorUtil.dotProduct(q, q) == 0) {
- continue;
- }
- for (int j = 0; j < nVectors; j++) {
- var v = vectors.get(j);
- dummyAccumulator += vsf.compare(q, v);
- }
+ double ashDummy = 0.0;
+ for (ForkJoinTask t : ashTasks) {
+ ashDummy += t.join();
+ }
+
+ long ashEnd = System.nanoTime();
+
+ System.out.println("\tASH dot-product computations took "
+ + (ashEnd - ashStart) / 1e9 + " seconds");
+
+ // ==================================================================
+ // [3] Float dot-product timing (ground truth baseline)
+ // ==================================================================
+ List> floatTasks = new ArrayList<>();
+
+ long floatStart = System.nanoTime();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ floatTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
+
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ for (int j = 0; j < nVectors; j++) {
+ localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
+ }
+ }
+ return localSum;
+ }));
+ }
+
+ double floatDummy = 0.0;
+ for (ForkJoinTask t : floatTasks) {
+ floatDummy += t.join();
}
- endTime = System.nanoTime();
- duration = (double) (endTime - startTime) / 1_000_000_000;
- System.out.println("\tFloat Distance computations took " + duration + " seconds");
- System.out.println("\tdummyAccumulator: " + dummyAccumulator);
+ long floatEnd = System.nanoTime();
+
+ System.out.println("\tFloat dot-product computations took "
+ + (floatEnd - floatStart) / 1e9 + " seconds");
+
+ // Prevent dead-code elimination
+ System.out.println("\tdummyAccumulator = " + (float) (ashDummy + floatDummy));
System.out.println("--");
}
@@ -117,7 +242,7 @@ public static void runSIFT() throws IOException {
var baseVectors = "siftsmall/siftsmall_base.fvecs";
var queryVectors = "siftsmall/siftsmall_query.fvecs";
- testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ testASHEncodings(baseVectors, queryVectors);
}
public static void runADA() throws IOException {
@@ -125,7 +250,7 @@ public static void runADA() throws IOException {
var baseVectors = "./fvec/wikipedia_squad/100k/ada_002_100000_base_vectors.fvec";
var queryVectors = "./fvec/wikipedia_squad/100k/ada_002_100000_query_vectors_10000.fvec";
- testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ testASHEncodings(baseVectors, queryVectors);
}
public static void runColbert() throws IOException {
@@ -133,7 +258,7 @@ public static void runColbert() throws IOException {
var baseVectors = "./fvec/wikipedia_squad/1M/colbertv2.0_128_base_vectors_1000000.fvec";
var queryVectors = "./fvec/wikipedia_squad/1M/colbertv2.0_128_query_vectors_100000.fvec";
- testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ testASHEncodings(baseVectors, queryVectors);
}
public static void runOpenai3072() throws IOException {
@@ -141,12 +266,12 @@ public static void runOpenai3072() throws IOException {
var baseVectors = "./fvec/wikipedia_squad/100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
var queryVectors = "./fvec/wikipedia_squad/100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
- testASHEncodings(baseVectors, queryVectors, VectorSimilarityFunction.COSINE);
+ testASHEncodings(baseVectors, queryVectors);
}
public static void main(String[] args) throws IOException {
runSIFT();
- runADA();
+// runADA();
runColbert();
// runOpenai3072();
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
index 5ada6365b..01d28c22f 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
@@ -7,17 +7,14 @@
public class TestAshRunWithoutTraining {
public static void main(String[] args) {
- int d = 1024; // input dimension
- int nDims = 128; // target dimension
-
- // Dummy data matrix X only to scrape shape
- RealMatrix X = MatrixUtils.createRealMatrix(10, d); // (10 × 1024)
+ int originalDim = 1024;
+ int quantizedDim = 128;
Random rng = new Random(42);
// Method under test
AsymmetricHashing.StiefelTransform stiefelTransform =
- AsymmetricHashing.runWithoutTraining(X, nDims, rng);
+ AsymmetricHashing.runWithoutTraining(originalDim, quantizedDim, rng);
// Inspect shapes and a few entries
RealMatrix A = stiefelTransform.A;
From 0dc6fb2b7e2e402ebe0c3c521966f9ba56a9a37e Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Fri, 26 Dec 2025 09:08:47 -0800
Subject: [PATCH 003/243] Enable float encoding and scoring of A, maintaining
double for training. Minor dataset updates as well.
---
.../jvector/quantization/ASHScorer.java | 8 ++--
.../quantization/AsymmetricHashing.java | 44 ++++++++++++++-----
.../jbellis/jvector/example/DistancesASH.java | 41 ++++++++++++++---
3 files changed, 70 insertions(+), 23 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
index 93f35182f..cff001511 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
@@ -53,7 +53,7 @@ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
// (avoid allocation by doing subtraction on the fly during projection)
final VectorFloat> mu = ash.globalMean;
- final double[][] A = ash.stiefelTransform.AData; // shape: d × originalDim
+ final float[][] A = ash.stiefelTransform.AFloat; // shape: d × originalDim
final int originalDim = ash.stiefelTransform.cols;
final float[] muArr = new float[originalDim];
for (int k = 0; k < originalDim; k++) muArr[k] = mu.get(k);
@@ -64,12 +64,12 @@ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
// Project query: tildeQ[j] = (A (q - μ))[j]
final float[] tildeQ = new float[d];
for (int j = 0; j < d; j++) {
- final double[] Aj = A[j];
- double acc = 0.0;
+ final float[] Aj = A[j];
+ float acc = 0.0f;
for (int k = 0; k < originalDim; k++) {
acc += Aj[k] * (qArr[k] - muArr[k]);
}
- tildeQ[j] = (float) acc;
+ tildeQ[j] = acc;
}
// = sum_j tildeQ[j]
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index a69fc9443..d3bfd0a38 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -321,8 +321,8 @@ public void quantizeBody(VectorFloat> vector,
StiefelTransform stiefel,
long[] outWords) {
- final double[][] A = stiefel.AData; // [quantizedDim][originalDim]
- final int originalDim = A[0].length;
+ final float[][] A = stiefel.AFloat;
+ final int originalDim = stiefel.cols;
// Copy vector and mean once (no virtual calls in inner loop)
final float[] x = new float[originalDim];
@@ -333,8 +333,8 @@ public void quantizeBody(VectorFloat> vector,
mu[d] = globalMean.get(d);
}
- // Eq. 6 normalization
- final double invNorm = (residualNorm > 0f) ? (1.0 / residualNorm) : 0.0;
+ // ASH paper, Eq. 6, normalization
+ final float invNorm = (residualNorm > 0f) ? (1.0f / residualNorm) : 0.0f;
final int words = QuantizedVector.wordsForDims(quantizedDim);
@@ -347,14 +347,14 @@ public void quantizeBody(VectorFloat> vector,
int bitIndex = base + j;
// y_j = dot(A[j], (x − μ) / ||x − μ||)
- double acc = 0.0;
- final double[] Arow = A[bitIndex];
+ float acc = 0.0f;
+ final float[] Arow = A[bitIndex];
for (int d = 0; d < originalDim; d++) {
acc += Arow[d] * (x[d] - mu[d]) * invNorm;
}
- if (acc > 0.0) {
+ if (acc > 0.0f) {
bits |= (1L << j);
}
}
@@ -559,6 +559,7 @@ private static ByteSequence> createBodyByteSequenceForBits(int nBits) {
* y = A · (x − μ)
*/
public static final class StiefelTransform {
+
/** Forward projection (column-orthonormal). */
public final RealMatrix W;
@@ -566,16 +567,25 @@ public static final class StiefelTransform {
public final RealMatrix A;
/**
- * Cached row-major backing for A when available.
- * This exists purely to avoid RealMatrix.getEntry() in hot loops.
+ * Cached row-major backing for A in double precision.
+ * Used for training, reference, and correctness-sensitive math.
+ *
* AData must not be mutated after construction.
*/
public final double[][] AData;
- /** Number of rows in A (i.e., quantizedDim = d). */
+ /**
+ * Cached row-major backing for A in float precision.
+ * Used for encoding and scoring fast paths.
+ *
+ * This is derived once from AData and must not be mutated.
+ */
+ public final float[][] AFloat;
+
+ /** Number of rows in A (quantizedDim = d). */
public final int rows;
- /** Number of columns in A (i.e., originalDim). */
+ /** Number of columns in A (originalDim). */
public final int cols;
public StiefelTransform(RealMatrix W) {
@@ -592,7 +602,17 @@ public StiefelTransform(RealMatrix W) {
}
this.rows = AData.length;
- this.cols = AData.length == 0 ? 0 : AData[0].length;
+ this.cols = rows == 0 ? 0 : AData[0].length;
+
+ // Build float copy once for fast encoding / scoring paths
+ this.AFloat = new float[rows][cols];
+ for (int i = 0; i < rows; i++) {
+ double[] src = AData[i];
+ float[] dst = AFloat[i];
+ for (int j = 0; j < cols; j++) {
+ dst[j] = (float) src[j];
+ }
+ }
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 695dc21e0..b56368cfe 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -80,7 +80,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = dimension; // sweep later if desired
+ int encodedBits = dimension / 4; // sweep later if desired
System.out.println(
"\toriginalDim = " + dimension +
@@ -245,6 +245,22 @@ public static void runSIFT() throws IOException {
testASHEncodings(baseVectors, queryVectors);
}
+ public static void runGIST() throws IOException {
+ System.out.println("Running GIST");
+
+ var baseVectors = "./fvec/gist/gist_base.fvecs";
+ var queryVectors = "./fvec/gist/gist_query.fvecs";
+ testASHEncodings(baseVectors, queryVectors);
+ }
+
+ public static void runCohere100k() throws IOException {
+ System.out.println("Running Cohere-100k");
+
+ var baseVectors = "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_base_vectors_100000.fvec";
+ var queryVectors = "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_query_vectors_10000.fvec";
+ testASHEncodings(baseVectors, queryVectors);
+ }
+
public static void runADA() throws IOException {
System.out.println("Running ada_002");
@@ -261,18 +277,29 @@ public static void runColbert() throws IOException {
testASHEncodings(baseVectors, queryVectors);
}
+ public static void runOpenai1536() throws IOException {
+ System.out.println("Running text-embedding-3-large_1536");
+
+ var baseVectors = "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_base_vectors.fvec";
+ var queryVectors = "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_query_vectors_10000.fvec";
+ testASHEncodings(baseVectors, queryVectors);
+ }
+
public static void runOpenai3072() throws IOException {
System.out.println("Running text-embedding-3-large_3072");
- var baseVectors = "./fvec/wikipedia_squad/100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
- var queryVectors = "./fvec/wikipedia_squad/100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
+ var baseVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
+ var queryVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
testASHEncodings(baseVectors, queryVectors);
}
public static void main(String[] args) throws IOException {
- runSIFT();
-// runADA();
- runColbert();
-// runOpenai3072();
+// runSIFT();
+// runGIST();
+// runColbert();
+ runCohere100k();
+ runADA();
+ runOpenai1536();
+ runOpenai3072();
}
}
\ No newline at end of file
From f730faad0a05a1589683ef90a18694588d1c6057 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Fri, 26 Dec 2025 11:08:10 -0800
Subject: [PATCH 004/243] Fix for shape of thin QR matrix. Plus scaffolding
for default (scalar) ASH backend. Validated ASH error and performance
correctness.
---
.../quantization/AsymmetricHashing.java | 147 ++++++++++++------
.../jvector/quantization/ash/AshMath.java | 46 ++++++
.../quantization/ash/DefaultAshMath.java | 70 +++++++++
.../example/TestAshRunWithoutTraining.java | 22 ++-
4 files changed, 227 insertions(+), 58 deletions(-)
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.java
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index d3bfd0a38..9efeb7e88 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -16,11 +16,14 @@
package io.github.jbellis.jvector.quantization;
+import io.github.jbellis.jvector.quantization.ash.AshMath;
+import io.github.jbellis.jvector.quantization.ash.DefaultAshMath;
import io.github.jbellis.jvector.annotations.VisibleForTesting;
import io.github.jbellis.jvector.disk.RandomAccessReader;
import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
import io.github.jbellis.jvector.util.Accountable;
+import io.github.jbellis.jvector.util.RamUsageEstimator;
import io.github.jbellis.jvector.vector.VectorUtil;
import io.github.jbellis.jvector.vector.VectorizationProvider;
import io.github.jbellis.jvector.vector.types.ByteSequence;
@@ -60,6 +63,8 @@ public class AsymmetricHashing implements VectorCompressor vector,
// Copy vector and mean once (no virtual calls in inner loop)
final float[] x = new float[originalDim];
final float[] mu = new float[originalDim];
-
for (int d = 0; d < originalDim; d++) {
x[d] = vector.get(d);
mu[d] = globalMean.get(d);
}
- // ASH paper, Eq. 6, normalization
+ // ASH paper Eq. 6 normalization
final float invNorm = (residualNorm > 0f) ? (1.0f / residualNorm) : 0.0f;
- final int words = QuantizedVector.wordsForDims(quantizedDim);
+ // Build normalized residual x̂ = (x − μ) / ||x − μ||
+ final float[] xhat = new float[originalDim];
+ for (int d = 0; d < originalDim; d++) {
+ xhat[d] = (x[d] - mu[d]) * invNorm;
+ }
+ // Project once: y = A · x̂
+ final float[] y = new float[quantizedDim];
+ ASH_MATH.project(A, xhat, y);
+
+ // Binarize y into outWords
+ final int words = QuantizedVector.wordsForDims(quantizedDim);
for (int w = 0; w < words; w++) {
long bits = 0L;
int base = w << 6;
int rem = Math.min(64, quantizedDim - base);
for (int j = 0; j < rem; j++) {
- int bitIndex = base + j;
-
- // y_j = dot(A[j], (x − μ) / ||x − μ||)
- float acc = 0.0f;
- final float[] Arow = A[bitIndex];
-
- for (int d = 0; d < originalDim; d++) {
- acc += Arow[d] * (x[d] - mu[d]) * invNorm;
- }
-
- if (acc > 0.0f) {
+ if (y[base + j] > 0.0f) {
bits |= (1L << j);
}
}
-
outWords[w] = bits;
}
}
@@ -560,71 +563,86 @@ private static ByteSequence> createBodyByteSequenceForBits(int nBits) {
*/
public static final class StiefelTransform {
- /** Forward projection (column-orthonormal). */
+ /** Forward projection (column-orthonormal). Shape: [originalDim × quantizedDim]. */
public final RealMatrix W;
- /** Transpose of W, used for projection. */
- public final RealMatrix A;
-
/**
- * Cached row-major backing for A in double precision.
+ * Cached row-major backing for A in double precision (A = Wᵀ).
* Used for training, reference, and correctness-sensitive math.
*
* AData must not be mutated after construction.
+ * Shape: [quantizedDim × originalDim].
*/
public final double[][] AData;
/**
- * Cached row-major backing for A in float precision.
+ * Cached row-major backing for A in float precision (A = Wᵀ).
* Used for encoding and scoring fast paths.
- *
- * This is derived once from AData and must not be mutated.
+ * AFloat must not be mutated after construction.
+ * Shape: [quantizedDim × originalDim].
*/
public final float[][] AFloat;
/** Number of rows in A (quantizedDim = d). */
public final int rows;
- /** Number of columns in A (originalDim). */
+ /** Number of columns in A (originalDim = D). */
public final int cols;
public StiefelTransform(RealMatrix W) {
- this.W = W;
- this.A = W.transpose();
-
- // Fast path: QRDecomposition + MatrixUtils typically produce Array2DRowRealMatrix
- // Use getDataRef() to avoid copying and avoid getEntry() calls.
- if (this.A instanceof Array2DRowRealMatrix) {
- this.AData = ((Array2DRowRealMatrix) this.A).getDataRef();
- } else {
- // Fallback: copy once at construction time (still far faster than getEntry per access)
- this.AData = this.A.getData();
+ this.W = Objects.requireNonNull(W, "W");
+
+ // Dimensions: W is [D × d]
+ final int originalDim = W.getRowDimension(); // D
+ final int quantizedDim = W.getColumnDimension(); // d
+
+ this.rows = quantizedDim;
+ this.cols = originalDim;
+
+ // Explicitly materialize A = Wᵀ into row-major arrays.
+ // We do this directly instead of relying on RealMatrix.transpose()+getData(),
+ // because some RealMatrix implementations return views where getData()
+ // can be surprisingly expensive or shaped unexpectedly.
+ this.AData = new double[quantizedDim][originalDim];
+ for (int i = 0; i < quantizedDim; i++) {
+ final double[] row = this.AData[i];
+ for (int j = 0; j < originalDim; j++) {
+ // A[i][j] = W[j][i]
+ row[j] = W.getEntry(j, i);
+ }
}
- this.rows = AData.length;
- this.cols = rows == 0 ? 0 : AData[0].length;
-
- // Build float copy once for fast encoding / scoring paths
- this.AFloat = new float[rows][cols];
- for (int i = 0; i < rows; i++) {
- double[] src = AData[i];
- float[] dst = AFloat[i];
- for (int j = 0; j < cols; j++) {
+ this.AFloat = new float[quantizedDim][originalDim];
+ for (int i = 0; i < quantizedDim; i++) {
+ final double[] src = this.AData[i];
+ final float[] dst = this.AFloat[i];
+ for (int j = 0; j < originalDim; j++) {
dst[j] = (float) src[j];
}
}
+
+ // Hard invariants: these are programmer-contract checks, not user input checks.
+ if (AFloat.length != rows) {
+ throw new AssertionError("AFloat rows " + AFloat.length + " != quantizedDim " + rows);
+ }
+ if (rows > 0 && AFloat[0].length != cols) {
+ throw new AssertionError("AFloat cols " + AFloat[0].length + " != originalDim " + cols);
+ }
+ if (AData.length != rows) {
+ throw new AssertionError("AData rows " + AData.length + " != quantizedDim " + rows);
+ }
+ if (rows > 0 && AData[0].length != cols) {
+ throw new AssertionError("AData cols " + AData[0].length + " != originalDim " + cols);
+ }
}
}
/**
* RANDOM (untrained) Stiefel initialization.
- *
* Generates a random orthonormal projection from R^originalDim → R^quantizedDim
* by sampling a Gaussian matrix and orthonormalizing via QR decomposition.
- *
* A deterministic sign normalization is applied so results are reproducible
* across runs and linear algebra backends.
- *
* @param originalDim dimensionality of the input vectors
* @param quantizedDim dimensionality of the binary embedding
* @param rng random number generator
@@ -650,13 +668,17 @@ public static StiefelTransform runWithoutTraining(int originalDim,
RealMatrix G = MatrixUtils.createRealMatrix(gaussian);
- // Thin QR decomposition: G = Q R
+ // QR decomposition: G = Q R
QRDecomposition qr = new QRDecomposition(G);
- RealMatrix Q = qr.getQ(); // originalDim × quantizedDim
- RealMatrix R = qr.getR();
+
+ // Commons Math returns full Q (m×m). We need thin Q (m×n).
+ RealMatrix Qfull = qr.getQ(); // originalDim × originalDim
+ RealMatrix Q = Qfull.getSubMatrix(0, originalDim - 1, 0, quantizedDim - 1); // originalDim × quantizedDim
+
+ RealMatrix R = qr.getR(); // originalDim × quantizedDim
// Normalize QR sign ambiguity for determinism (RANDOM only)
- int diag = Math.min(R.getRowDimension(), R.getColumnDimension());
+ int diag = Math.min(R.getRowDimension(), R.getColumnDimension()); // == quantizedDim
for (int j = 0; j < diag; j++) {
if (R.getEntry(j, j) < 0.0) {
for (int i = 0; i < Q.getRowDimension(); i++) {
@@ -665,6 +687,7 @@ public static StiefelTransform runWithoutTraining(int originalDim,
}
}
+ // W := Q (originalDim × quantizedDim)
return new StiefelTransform(Q);
}
@@ -840,10 +863,32 @@ public long ramBytesUsed() {
}
if (stiefelTransform != null) {
+ // Authoritative double-precision model
bytes += estimateMatrixBytes(stiefelTransform.W);
- bytes += estimateMatrixBytes(stiefelTransform.A);
+
+ // Materialized compute buffers (row-major)
+ bytes += estimateArray2DBytes(stiefelTransform.AData);
+ bytes += estimateArray2DBytes(stiefelTransform.AFloat);
+ }
+
+ return bytes;
+ }
+
+ private static long estimateArray2DBytes(double[][] a) {
+ if (a == null) return 0L;
+ long bytes = 0L;
+ for (double[] row : a) {
+ bytes += RamUsageEstimator.sizeOf(row);
}
+ return bytes;
+ }
+ private static long estimateArray2DBytes(float[][] a) {
+ if (a == null) return 0L;
+ long bytes = 0L;
+ for (float[] row : a) {
+ bytes += RamUsageEstimator.sizeOf(row);
+ }
return bytes;
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java
new file mode 100644
index 000000000..8d277f576
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java
@@ -0,0 +1,46 @@
+/*
+ * 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.ash;
+
+/**
+ * ASH math kernels.
+ *
+ * This interface forms the SIMD seam for ASH-specific operations.
+ * Implementations may be scalar, Panama Vector API, or native.
+ */
+public interface AshMath {
+
+ /**
+ * Compute y = A · x, where:
+ * - A is [d][D]
+ * - x is [D]
+ * - y is [d]
+ */
+ void project(float[][] A, float[] x, float[] y);
+
+ /**
+ * Compute masked-add ⟨tildeQ, b⟩ where b ∈ {0,1}^d stored as bitpacked longs.
+ * Implementations must interpret bits with the same convention as existing ASH code:
+ * bit j corresponds to dimension j (little-endian within each long).
+ *
+ * @param tildeQ projected query (length d)
+ * @param bits bitpacked database code
+ * @param d number of valid dimensions in tildeQ
+ * @return sum_{j: b_j=1} tildeQ[j]
+ */
+ float maskedAdd(float[] tildeQ, long[] bits, int d);
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.java
new file mode 100644
index 000000000..a177f7fe5
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.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.quantization.ash;
+
+/**
+ * Scalar reference implementation of ASH math kernels.
+ */
+public final class DefaultAshMath implements AshMath {
+
+ @Override
+ public void project(float[][] A, float[] x, float[] y) {
+ final int rows = A.length;
+ final int cols = (rows == 0 ? 0 : A[0].length);
+
+ assert y.length == rows :
+ "ASH project shape mismatch: y.length=" + y.length +
+ ", A.rows=" + rows +
+ ", A.cols=" + cols +
+ ", x.length=" + x.length;
+
+ assert cols == x.length :
+ "ASH project shape mismatch: A.cols=" + cols +
+ " != x.length=" + x.length +
+ " (A.rows=" + rows +
+ ", y.length=" + y.length + ")";
+
+ for (int i = 0; i < rows; i++) {
+ float acc = 0.0f;
+ final float[] Arow = A[i];
+ for (int j = 0; j < cols; j++) {
+ acc += Arow[j] * x[j];
+ }
+ y[i] = acc;
+ }
+ }
+
+ @Override
+ public float maskedAdd(float[] tildeQ, long[] bits, int d) {
+ assert d <= tildeQ.length : "d must be <= tildeQ.length";
+
+ float sum = 0.0f;
+ int base = 0;
+
+ for (int w = 0; w < bits.length && base < d; w++, base += 64) {
+ long word = bits[w];
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = base + bit;
+ if (idx < d) sum += tildeQ[idx];
+ word &= (word - 1);
+ }
+ }
+
+ return sum;
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
index 01d28c22f..9e6ee5566 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/TestAshRunWithoutTraining.java
@@ -1,7 +1,9 @@
package io.github.jbellis.jvector.example;
import io.github.jbellis.jvector.quantization.AsymmetricHashing;
-import org.apache.commons.math3.linear.*;
+import org.apache.commons.math3.linear.MatrixUtils;
+import org.apache.commons.math3.linear.RealMatrix;
+
import java.util.Random;
public class TestAshRunWithoutTraining {
@@ -16,14 +18,20 @@ public static void main(String[] args) {
AsymmetricHashing.StiefelTransform stiefelTransform =
AsymmetricHashing.runWithoutTraining(originalDim, quantizedDim, rng);
- // Inspect shapes and a few entries
- RealMatrix A = stiefelTransform.A;
- RealMatrix W = stiefelTransform.W;
+ // Shapes
+ System.out.println("A shape = " +
+ stiefelTransform.rows + " x " + stiefelTransform.cols);
+ System.out.println("W shape = " +
+ stiefelTransform.W.getRowDimension() + " x " +
+ stiefelTransform.W.getColumnDimension());
- System.out.println("A shape = " + A.getRowDimension() + " x " + A.getColumnDimension());
- System.out.println("W shape = " + W.getRowDimension() + " x " + W.getColumnDimension());
+ // Reconstruct A as a RealMatrix *for testing only*
+ RealMatrix A =
+ MatrixUtils.createRealMatrix(stiefelTransform.AData);
+
+ RealMatrix W = stiefelTransform.W;
- // Sanity check....
+ // Sanity check: a few entries of A
System.out.println("First few entries of A:");
for (int i = 0; i < Math.min(3, A.getRowDimension()); i++) {
for (int j = 0; j < Math.min(3, A.getColumnDimension()); j++) {
From e91a842e837f9eb8b8f8dc7402910d19496b15ad Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Fri, 26 Dec 2025 19:49:17 -0800
Subject: [PATCH 005/243] First implementation with fast scalar and SIMD
projection encoding.
---
.../jvector/quantization/ASHScorer.java | 6 +
.../quantization/AsymmetricHashing.java | 23 ++-
.../jvector/quantization/ash/AshMath.java | 46 ------
.../quantization/ash/DefaultAshMath.java | 70 --------
.../jvector/vector/VectorUtilSupport.java | 43 +++++
.../vector/PanamaVectorUtilSupport.java | 151 +++++++++++-------
6 files changed, 149 insertions(+), 190 deletions(-)
delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java
delete mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.java
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
index cff001511..f3a2cfea9 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
@@ -56,6 +56,12 @@ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
final float[][] A = ash.stiefelTransform.AFloat; // shape: d × originalDim
final int originalDim = ash.stiefelTransform.cols;
final float[] muArr = new float[originalDim];
+
+ final var vecUtil =
+ io.github.jbellis.jvector.vector.VectorizationProvider
+ .getInstance()
+ .getVectorUtilSupport();
+
for (int k = 0; k < originalDim; k++) muArr[k] = mu.get(k);
final float[] qArr = new float[originalDim];
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 9efeb7e88..9f5ca260a 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -16,8 +16,6 @@
package io.github.jbellis.jvector.quantization;
-import io.github.jbellis.jvector.quantization.ash.AshMath;
-import io.github.jbellis.jvector.quantization.ash.DefaultAshMath;
import io.github.jbellis.jvector.annotations.VisibleForTesting;
import io.github.jbellis.jvector.disk.RandomAccessReader;
import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
@@ -63,8 +61,6 @@ public class AsymmetricHashing implements VectorCompressor vector,
final float[][] A = stiefel.AFloat;
final int originalDim = stiefel.cols;
+ final var vecUtil =
+ io.github.jbellis.jvector.vector.VectorizationProvider
+ .getInstance()
+ .getVectorUtilSupport();
+
// Copy vector and mean once (no virtual calls in inner loop)
final float[] x = new float[originalDim];
final float[] mu = new float[originalDim];
@@ -337,20 +338,16 @@ public void quantizeBody(VectorFloat> vector,
mu[d] = globalMean.get(d);
}
- // ASH paper Eq. 6 normalization
+ // ASH paper, Eq. 6, normalization
final float invNorm = (residualNorm > 0f) ? (1.0f / residualNorm) : 0.0f;
- // Build normalized residual x̂ = (x − μ) / ||x − μ||
+ // Normalized residual x̂
final float[] xhat = new float[originalDim];
for (int d = 0; d < originalDim; d++) {
xhat[d] = (x[d] - mu[d]) * invNorm;
}
- // Project once: y = A · x̂
- final float[] y = new float[quantizedDim];
- ASH_MATH.project(A, xhat, y);
-
- // Binarize y into outWords
+ // Binarize directly from per-row projection
final int words = QuantizedVector.wordsForDims(quantizedDim);
for (int w = 0; w < words; w++) {
long bits = 0L;
@@ -358,7 +355,9 @@ public void quantizeBody(VectorFloat> vector,
int rem = Math.min(64, quantizedDim - base);
for (int j = 0; j < rem; j++) {
- if (y[base + j] > 0.0f) {
+ int bitIndex = base + j;
+ float acc = vecUtil.ashDotRow(A[bitIndex], xhat);
+ if (acc > 0.0f) {
bits |= (1L << j);
}
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java
deleted file mode 100644
index 8d277f576..000000000
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/AshMath.java
+++ /dev/null
@@ -1,46 +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.
- */
-
-package io.github.jbellis.jvector.quantization.ash;
-
-/**
- * ASH math kernels.
- *
- * This interface forms the SIMD seam for ASH-specific operations.
- * Implementations may be scalar, Panama Vector API, or native.
- */
-public interface AshMath {
-
- /**
- * Compute y = A · x, where:
- * - A is [d][D]
- * - x is [D]
- * - y is [d]
- */
- void project(float[][] A, float[] x, float[] y);
-
- /**
- * Compute masked-add ⟨tildeQ, b⟩ where b ∈ {0,1}^d stored as bitpacked longs.
- * Implementations must interpret bits with the same convention as existing ASH code:
- * bit j corresponds to dimension j (little-endian within each long).
- *
- * @param tildeQ projected query (length d)
- * @param bits bitpacked database code
- * @param d number of valid dimensions in tildeQ
- * @return sum_{j: b_j=1} tildeQ[j]
- */
- float maskedAdd(float[] tildeQ, long[] bits, int d);
-}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.java
deleted file mode 100644
index a177f7fe5..000000000
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ash/DefaultAshMath.java
+++ /dev/null
@@ -1,70 +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.
- */
-
-package io.github.jbellis.jvector.quantization.ash;
-
-/**
- * Scalar reference implementation of ASH math kernels.
- */
-public final class DefaultAshMath implements AshMath {
-
- @Override
- public void project(float[][] A, float[] x, float[] y) {
- final int rows = A.length;
- final int cols = (rows == 0 ? 0 : A[0].length);
-
- assert y.length == rows :
- "ASH project shape mismatch: y.length=" + y.length +
- ", A.rows=" + rows +
- ", A.cols=" + cols +
- ", x.length=" + x.length;
-
- assert cols == x.length :
- "ASH project shape mismatch: A.cols=" + cols +
- " != x.length=" + x.length +
- " (A.rows=" + rows +
- ", y.length=" + y.length + ")";
-
- for (int i = 0; i < rows; i++) {
- float acc = 0.0f;
- final float[] Arow = A[i];
- for (int j = 0; j < cols; j++) {
- acc += Arow[j] * x[j];
- }
- y[i] = acc;
- }
- }
-
- @Override
- public float maskedAdd(float[] tildeQ, long[] bits, int d) {
- assert d <= tildeQ.length : "d must be <= tildeQ.length";
-
- float sum = 0.0f;
- int base = 0;
-
- for (int w = 0; w < bits.length && base < d; w++, base += 64) {
- long word = bits[w];
- while (word != 0L) {
- int bit = Long.numberOfTrailingZeros(word);
- int idx = base + bit;
- if (idx < d) sum += tildeQ[idx];
- word &= (word - 1);
- }
- }
-
- return sum;
- }
-}
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 d8223ab12..417d513b1 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
@@ -235,4 +235,47 @@ default float pqDecodedCosineSimilarity(ByteSequence> encoded, int encodedOffs
*/
float nvqUniformLoss(VectorFloat> vector, float minValue, float maxValue, int nBits);
+ // ------------------------------------------------------------------
+ // ASH kernels (default scalar implementations)
+ // ------------------------------------------------------------------
+
+ /**
+ * Dot product between a matrix row and a dense vector.
+ * Used by ASH projection: y_i = .
+ *
+ * Contract: Arow.length == x.length.
+ * Default is scalar; SIMD backends may override.
+ */
+ default float ashDotRow(float[] Arow, float[] x) {
+ assert Arow.length == x.length : "Arow.length != x.length";
+ float acc = 0.0f;
+ for (int i = 0; i < x.length; i++) {
+ acc += Arow[i] * x[i];
+ }
+ return acc;
+ }
+
+ /**
+ * maskedAdd = for b ∈ {0,1}^d stored as bitpacked longs.
+ *
+ * Contract: d <= tildeQ.length.
+ * Default is scalar; SIMD backends may override.
+ */
+ default float ashMaskedAdd(float[] tildeQ, long[] bits, int d) {
+ assert d <= tildeQ.length : "d must be <= tildeQ.length";
+
+ float sum = 0.0f;
+ int base = 0;
+
+ for (int w = 0; w < bits.length && base < d; w++, base += 64) {
+ long word = bits[w];
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = base + bit;
+ if (idx < d) sum += tildeQ[idx];
+ word &= (word - 1);
+ }
+ }
+ return sum;
+ }
}
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 dc18c4bf2..ea2211e8f 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
@@ -43,7 +43,7 @@ protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat>
}
protected FloatVector fromVectorFloat(VectorSpecies SPEC, VectorFloat> vector, int offset, int[] indices, int indicesOffset) {
- return FloatVector.fromArray(SPEC, ((ArrayVectorFloat)vector).get(), offset, indices, indicesOffset);
+ return FloatVector.fromArray(SPEC, ((ArrayVectorFloat) vector).get(), offset, indices, indicesOffset);
}
protected void intoVectorFloat(FloatVector vector, VectorFloat> v, int offset) {
@@ -148,8 +148,7 @@ public float dotProduct(VectorFloat> v1, VectorFloat> v2) {
}
@Override
- public float dotProduct(VectorFloat> v1, int v1offset, VectorFloat> v2, int v2offset, final int length)
- {
+ public float dotProduct(VectorFloat> v1, int v1offset, VectorFloat> v2, int v2offset, final int length) {
//Common case first
if (length >= FloatVector.SPECIES_PREFERRED.length())
return dotProductPreferred(v1, v1offset, v2, v2offset, length);
@@ -251,8 +250,7 @@ float dotProductPreferred(VectorFloat> va, int vaoffset, VectorFloat> vb, in
// Unrolled vector loop; for dot product from L1 cache, an unroll factor of 2 generally suffices.
// If we are going to be getting data that's further down the hierarchy but not fetched off disk/network,
// we might want to unroll further, e.g. to 8 (4 sets of a,b,sum with 3-ahead reads seems to work best).
- if (length >= vectorLength * 2)
- {
+ if (length >= vectorLength * 2) {
length -= vectorLength * 2;
a0 = fromVectorFloat(FloatVector.SPECIES_PREFERRED, va, vaoffset + vectorLength * 0);
b0 = fromVectorFloat(FloatVector.SPECIES_PREFERRED, vb, vboffset + vectorLength * 0);
@@ -260,8 +258,7 @@ float dotProductPreferred(VectorFloat> va, int vaoffset, VectorFloat> vb, in
b1 = fromVectorFloat(FloatVector.SPECIES_PREFERRED, vb, vboffset + vectorLength * 1);
vaoffset += vectorLength * 2;
vboffset += vectorLength * 2;
- while (length >= vectorLength * 2)
- {
+ while (length >= vectorLength * 2) {
// All instructions in the main loop have no dependencies between them and can be executed in parallel.
length -= vectorLength * 2;
sum0 = a0.fma(b0, sum0);
@@ -393,8 +390,7 @@ public float squareDistance(VectorFloat> v1, VectorFloat> v2) {
}
@Override
- public float squareDistance(VectorFloat> v1, int v1offset, VectorFloat> v2, int v2offset, final int length)
- {
+ public float squareDistance(VectorFloat> v1, int v1offset, VectorFloat> v2, int v2offset, final int length) {
//Common case first
if (length >= FloatVector.SPECIES_PREFERRED.length())
return squareDistancePreferred(v1, v1offset, v2, v2offset, length);
@@ -550,7 +546,7 @@ public void addInPlace(VectorFloat> v1, VectorFloat> v2) {
// Process the tail
for (int i = vectorizedLength; i < v1.length(); i++) {
- v1.set(i, v1.get(i) + v2.get(i));
+ v1.set(i, v1.get(i) + v2.get(i));
}
}
@@ -571,7 +567,7 @@ public void addInPlace(VectorFloat> v1, float value) {
// Process the tail
for (int i = vectorizedLength; i < v1.length(); i++) {
- v1.set(i, v1.get(i) + value);
+ v1.set(i, v1.get(i) + value);
}
}
@@ -592,7 +588,7 @@ public void subInPlace(VectorFloat> v1, VectorFloat> v2) {
// Process the tail
for (int i = vectorizedLength; i < v1.length(); i++) {
- v1.set(i, v1.get(i) - v2.get(i));
+ v1.set(i, v1.get(i) - v2.get(i));
}
}
@@ -608,7 +604,7 @@ public void subInPlace(VectorFloat> vector, float value) {
// Process the tail
for (int i = vectorizedLength; i < vector.length(); i++) {
- vector.set(i, vector.get(i) - value);
+ vector.set(i, vector.get(i) - value);
}
}
@@ -679,19 +675,18 @@ public void minInPlace(VectorFloat> v1, VectorFloat> v2) {
// Process the tail
for (int i = vectorizedLength; i < v1.length(); i++) {
- v1.set(i, Math.min(v1.get(i), v2.get(i)));
+ v1.set(i, Math.min(v1.get(i), v2.get(i)));
}
}
@Override
public float assembleAndSum(VectorFloat> data, int dataBase, ByteSequence> baseOffsets) {
- return assembleAndSum(data, dataBase, baseOffsets, 0, baseOffsets.length());
+ return assembleAndSum(data, dataBase, baseOffsets, 0, baseOffsets.length());
}
@Override
public float assembleAndSum(VectorFloat> data, int dataBase, ByteSequence> baseOffsets, int baseOffsetsOffset, int baseOffsetsLength) {
- return switch (PREFERRED_BIT_SIZE)
- {
+ return switch (PREFERRED_BIT_SIZE) {
case 512 -> assembleAndSum512(data, dataBase, baseOffsets, baseOffsetsOffset, baseOffsetsLength);
case 256 -> assembleAndSum256(data, dataBase, baseOffsets, baseOffsetsOffset, baseOffsetsLength);
case 128 -> assembleAndSum128(data, dataBase, baseOffsets, baseOffsetsOffset, baseOffsetsLength);
@@ -712,7 +707,7 @@ float assembleAndSum512(VectorFloat> data, int dataBase, ByteSequence> baseO
.lanewise(VectorOperators.AND, BYTE_TO_INT_MASK_512)
.reinterpretAsInts()
.add(scale)
- .intoArray(convOffsets,0);
+ .intoArray(convOffsets, 0);
var offset = i * dataBase;
sum = sum.add(fromVectorFloat(FloatVector.SPECIES_512, data, offset, convOffsets, 0));
@@ -741,7 +736,7 @@ float assembleAndSum256(VectorFloat> data, int dataBase, ByteSequence> baseO
.lanewise(VectorOperators.AND, BYTE_TO_INT_MASK_256)
.reinterpretAsInts()
.add(scale)
- .intoArray(convOffsets,0);
+ .intoArray(convOffsets, 0);
var offset = i * dataBase;
sum = sum.add(fromVectorFloat(FloatVector.SPECIES_256, data, offset, convOffsets, 0));
@@ -776,11 +771,13 @@ public float assembleAndSumPQ(
int clusterCount // = k
) {
//compute the size of the subvector
- return switch (PREFERRED_BIT_SIZE)
- {
- case 512 -> assembleAndSumPQ_512(codebookPartialSums, subspaceCount, vector1Ordinals, vector1OrdinalOffset, vector2Ordinals, vector2OrdinalOffset, clusterCount);
- case 256 -> assembleAndSumPQ_256(codebookPartialSums, subspaceCount, vector1Ordinals, vector1OrdinalOffset, vector2Ordinals, vector2OrdinalOffset, clusterCount);
- case 128 -> assembleAndSumPQ_128(codebookPartialSums, subspaceCount, vector1Ordinals, vector1OrdinalOffset, vector2Ordinals, vector2OrdinalOffset, clusterCount);
+ return switch (PREFERRED_BIT_SIZE) {
+ case 512 ->
+ assembleAndSumPQ_512(codebookPartialSums, subspaceCount, vector1Ordinals, vector1OrdinalOffset, vector2Ordinals, vector2OrdinalOffset, clusterCount);
+ case 256 ->
+ assembleAndSumPQ_256(codebookPartialSums, subspaceCount, vector1Ordinals, vector1OrdinalOffset, vector2Ordinals, vector2OrdinalOffset, clusterCount);
+ case 128 ->
+ assembleAndSumPQ_128(codebookPartialSums, subspaceCount, vector1Ordinals, vector1OrdinalOffset, vector2Ordinals, vector2OrdinalOffset, clusterCount);
default -> throw new IllegalStateException("Unsupported vector width: " + PREFERRED_BIT_SIZE);
};
}
@@ -794,19 +791,19 @@ float assembleAndSumPQ_128(
int baseOffsetsOffset2,
int clusterCount // = k
) {
- final int k = clusterCount;
- final int blockSize = k * (k + 1) / 2;
+ final int k = clusterCount;
+ final int blockSize = k * (k + 1) / 2;
float res = 0f;
for (int i = 0; i < subspaceCount; i++) {
int c1 = Byte.toUnsignedInt(baseOffsets1.get(i + baseOffsetsOffset1));
int c2 = Byte.toUnsignedInt(baseOffsets2.get(i + baseOffsetsOffset2));
- int r = Math.min(c1, c2);
- int c = Math.max(c1, c2);
+ int r = Math.min(c1, c2);
+ int c = Math.max(c1, c2);
- int offsetRow = r * k - (r * (r - 1) / 2);
+ int offsetRow = r * k - (r * (r - 1) / 2);
int idxInBlock = offsetRow + (c - r);
- int base = i * blockSize;
+ int base = i * blockSize;
res += data.get(base + idxInBlock);
}
@@ -824,10 +821,10 @@ float assembleAndSumPQ_256(
int clusterCount // = k
) {
final VectorSpecies FSPECIES = FloatVector.SPECIES_256;
- final int LANES = FSPECIES.length();
- final int k = clusterCount;
- final int blockSize = k * (k + 1) / 2;
- final int M = subspaceCount;
+ final int LANES = FSPECIES.length();
+ final int k = clusterCount;
+ final int blockSize = k * (k + 1) / 2;
+ final int M = subspaceCount;
int[] convOffsets = scratchInt256.get();
FloatVector sum = FloatVector.zero(FloatVector.SPECIES_256);
@@ -856,9 +853,9 @@ float assembleAndSumPQ_256(
var c = c1v.max(c2v);
// c) offsetRow = r*k - (r*(r-1))/2
- var rk = r.mul(kvec);
- var triangular = r.mul(r.sub(onevec)).mul(twovec);
- var offsetRow = rk.sub(triangular);
+ var rk = r.mul(kvec);
+ var triangular = r.mul(r.sub(onevec)).mul(twovec);
+ var offsetRow = rk.sub(triangular);
// d) idxInBlock = offsetRow + (c - r) + (i * blockSize)
offsetRow.add(c.sub(r)).add(scale)
@@ -881,12 +878,12 @@ float assembleAndSumPQ_256(
for (int i = (M / LANES) * LANES; i < M; i++) {
int c1 = Byte.toUnsignedInt(baseOffsets1.get(i + baseOffsetsOffset1));
int c2 = Byte.toUnsignedInt(baseOffsets2.get(i + baseOffsetsOffset2));
- int r = Math.min(c1, c2);
- int c = Math.max(c1, c2);
+ int r = Math.min(c1, c2);
+ int c = Math.max(c1, c2);
- int offsetRow = r * k - (r * (r - 1) / 2);
+ int offsetRow = r * k - (r * (r - 1) / 2);
int idxInBlock = offsetRow + (c - r);
- int base = i * blockSize;
+ int base = i * blockSize;
res += data.get(base + idxInBlock);
}
@@ -904,10 +901,10 @@ float assembleAndSumPQ_512(
int clusterCount // = k
) {
final VectorSpecies FSPECIES = FloatVector.SPECIES_512;
- final int LANES = FSPECIES.length();
- final int k = clusterCount;
- final int blockSize = k * (k + 1) / 2;
- final int M = subspaceCount;
+ final int LANES = FSPECIES.length();
+ final int k = clusterCount;
+ final int blockSize = k * (k + 1) / 2;
+ final int M = subspaceCount;
int[] convOffsets = scratchInt512.get();
FloatVector sum = FloatVector.zero(FloatVector.SPECIES_512);
@@ -934,9 +931,9 @@ float assembleAndSumPQ_512(
var c = c1v.max(c2v);
// c) offsetRow = r*k - (r*(r-1))/2
- var rk = r.mul(kvec);
- var triangular = r.mul(r.sub(onevec)).mul(twovec);
- var offsetRow = rk.sub(triangular);
+ var rk = r.mul(kvec);
+ var triangular = r.mul(r.sub(onevec)).mul(twovec);
+ var offsetRow = rk.sub(triangular);
// d) idxInBlock = offsetRow + (c - r) + (i * blockSize)
offsetRow.add(c.sub(r)).add(scale)
@@ -959,12 +956,12 @@ float assembleAndSumPQ_512(
for (int i = (M / LANES) * LANES; i < M; i++) {
int c1 = Byte.toUnsignedInt(baseOffsets1.get(i + baseOffsetsOffset1));
int c2 = Byte.toUnsignedInt(baseOffsets2.get(i + baseOffsetsOffset2));
- int r = Math.min(c1, c2);
- int c = Math.max(c1, c2);
+ int r = Math.min(c1, c2);
+ int c = Math.max(c1, c2);
- int offsetRow = r * k - (r * (r - 1) / 2);
+ int offsetRow = r * k - (r * (r - 1) / 2);
int idxInBlock = offsetRow + (c - r);
- int base = i * blockSize;
+ int base = i * blockSize;
res += data.get(base + idxInBlock);
}
@@ -1040,7 +1037,7 @@ private static int combineBytes(int i, int shuffle, ByteSequence> quantizedPar
return ((Byte.toUnsignedInt(highByte) << 8) | Byte.toUnsignedInt(lowByte));
}
- private static float combineBytes(int i, int shuffle, VectorFloat> partials) {
+ private static float combineBytes(int i, int shuffle, VectorFloat> partials) {
return partials.get(i * 256 + shuffle);
}
@@ -1056,9 +1053,12 @@ private static int computeSingleShuffle(int codebookPosition, int neighborPositi
@Override
public float pqDecodedCosineSimilarity(ByteSequence> encoded, int encodedOffset, int encodedLength, int clusterCount, VectorFloat> partialSums, VectorFloat> aMagnitude, float bMagnitude) {
return switch (PREFERRED_BIT_SIZE) {
- case 512 -> pqDecodedCosineSimilarity512(encoded, encodedOffset, encodedLength, clusterCount, partialSums, aMagnitude, bMagnitude);
- case 256 -> pqDecodedCosineSimilarity256(encoded, encodedOffset, encodedLength, clusterCount, partialSums, aMagnitude, bMagnitude);
- case 128 -> pqDecodedCosineSimilarity128(encoded, encodedOffset, encodedLength, clusterCount, partialSums, aMagnitude, bMagnitude);
+ case 512 ->
+ pqDecodedCosineSimilarity512(encoded, encodedOffset, encodedLength, clusterCount, partialSums, aMagnitude, bMagnitude);
+ case 256 ->
+ pqDecodedCosineSimilarity256(encoded, encodedOffset, encodedLength, clusterCount, partialSums, aMagnitude, bMagnitude);
+ case 128 ->
+ pqDecodedCosineSimilarity128(encoded, encodedOffset, encodedLength, clusterCount, partialSums, aMagnitude, bMagnitude);
default -> throw new IllegalStateException("Unsupported vector width: " + PREFERRED_BIT_SIZE);
};
}
@@ -1080,7 +1080,7 @@ float pqDecodedCosineSimilarity512(ByteSequence> baseOffsets, int baseOffsetsO
.lanewise(VectorOperators.AND, BYTE_TO_INT_MASK_512)
.reinterpretAsInts()
.add(scale)
- .intoArray(convOffsets,0);
+ .intoArray(convOffsets, 0);
var offset = i * clusterCount;
sum = sum.add(fromVectorFloat(FloatVector.SPECIES_512, partialSums, offset, convOffsets, 0));
@@ -1116,7 +1116,7 @@ float pqDecodedCosineSimilarity256(ByteSequence> baseOffsets, int baseOffsetsO
.lanewise(VectorOperators.AND, BYTE_TO_INT_MASK_256)
.reinterpretAsInts()
.add(scale)
- .intoArray(convOffsets,0);
+ .intoArray(convOffsets, 0);
var offset = i * clusterCount;
sum = sum.add(fromVectorFloat(FloatVector.SPECIES_256, partialSums, offset, convOffsets, 0));
@@ -1492,7 +1492,7 @@ public float[] nvqCosine8bit(VectorFloat> vector,
void transpose(VectorFloat> arr, int first, int last, int nRows) {
final int mn1 = (last - first - 1);
- final int n = (last - first) / nRows;
+ final int n = (last - first) / nRows;
boolean[] visited = new boolean[last - first];
float temp;
int cycle = first;
@@ -1500,7 +1500,7 @@ void transpose(VectorFloat> arr, int first, int last, int nRows) {
if (visited[cycle - first])
continue;
int a = cycle - first;
- do {
+ do {
a = a == mn1 ? mn1 : (n * a) % mn1;
temp = arr.get(first + a);
arr.set(first + a, arr.get(cycle));
@@ -1541,6 +1541,33 @@ public void calculatePartialSums(VectorFloat> codebook, int codebookIndex, int
@Override
public float pqDecodedCosineSimilarity(ByteSequence> encoded, int clusterCount, VectorFloat> partialSums, VectorFloat> aMagnitude, float bMagnitude) {
- return pqDecodedCosineSimilarity(encoded, 0, encoded.length(), clusterCount, partialSums, aMagnitude, bMagnitude);
+ return pqDecodedCosineSimilarity(encoded, 0, encoded.length(), clusterCount, partialSums, aMagnitude, bMagnitude);
+ }
+
+ @Override
+ public float ashDotRow(float[] Arow, float[] x) {
+ assert Arow.length == x.length : "Arow.length != x.length";
+
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int len = x.length;
+ final int upper = SPEC.loopBound(len);
+
+ FloatVector sum = FloatVector.zero(SPEC);
+
+ int i = 0;
+ for (; i < upper; i += SPEC.length()) {
+ FloatVector va = FloatVector.fromArray(SPEC, Arow, i);
+ FloatVector vx = FloatVector.fromArray(SPEC, x, i);
+ sum = va.fma(vx, sum);
+ }
+
+ float acc = sum.reduceLanes(VectorOperators.ADD);
+
+ // scalar tail
+ for (; i < len; i++) {
+ acc += Arow[i] * x[i];
+ }
+
+ return acc;
}
-}
\ No newline at end of file
+}
From d2b6b0e713d0e3dcb90ac44a343dee428f472569 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Tue, 30 Dec 2025 08:20:58 -0800
Subject: [PATCH 006/243] First implementation of paper-faithful block
column-major scorer with SIMD support.
---
.../jvector/quantization/ASHBlockScorer.java | 23 +
.../jvector/quantization/ASHVectors.java | 418 +++++++++++++++++-
.../jvector/vector/VectorUtilSupport.java | 57 ++-
.../jbellis/jvector/example/DistancesASH.java | 165 ++++++-
.../vector/PanamaVectorUtilSupport.java | 271 ++++++++++++
5 files changed, 902 insertions(+), 32 deletions(-)
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHBlockScorer.java
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHBlockScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHBlockScorer.java
new file mode 100644
index 000000000..f2555d748
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHBlockScorer.java
@@ -0,0 +1,23 @@
+package io.github.jbellis.jvector.quantization;
+
+/**
+ * Block-oriented scorer for ASH vectors.
+ * A block scorer computes scores for a contiguous range of vector ordinals
+ * into a caller-provided output buffer.
+ * This interface exists to support:
+ * - Blocked SIMD scoring
+ * - Query reuse across multiple neighbors
+ * - Graph-local scoring (e.g., FusedASH)
+ */
+public interface ASHBlockScorer {
+
+ /**
+ * Score {@code count} vectors starting at {@code start}.
+ * Results are written to {@code out[0..count-1]}.
+ * @param start starting ordinal
+ * @param count number of vectors to score
+ * @param out output buffer (must have length ≥ count)
+ */
+ void scoreRange(int start, int count, float[] out);
+}
+
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index dcfe63a39..5ee690c6e 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -32,22 +32,21 @@
* Container for ASH-compressed vectors with a minimal scoring interface.
*
*
- * This class intentionally implements only the functionality required for:
+ * This class implements the functionality required for:
*
- * - Encoding benchmarks
- * - Approximate distance microbenchmarks
- * - Correctness validation against full-precision distances
+ * - Scalar scoring
+ * - Block scoring
+ * - Vectorized versions of scoring
*
*
*
* TODO (future work):
*
- * - Implement landmark-aware ASH scoring using residualNorm and dotWithLandmark
* - Add a full ASH ScoreFunction aligned with graph search semantics
* - Support similarityToNeighbor(...) where appropriate
* - Integrate ASH as a Feature for on-disk graph indexes
- * - Define reranking strategy (likely via full vectors)
- * - Support multi-landmark ASH (this is intentionally deferred)
+ * - Define reranking strategy
+ * - Support multi-landmark ASH
*
*/
public class ASHVectors implements CompressedVectors {
@@ -56,6 +55,19 @@ public class ASHVectors implements CompressedVectors {
final AsymmetricHashing.QuantizedVector[] compressedVectors;
final ASHScorer scorer;
+ // Cached scalar headers (avoid object chasing in blocked scorer)
+ private final float[] residualNorms;
+ private final float[] dotXMu;
+
+ // Quantized dimensions
+ private final int d;
+ private final int words;
+
+ // Packed block-column-major bits (built lazily per blockSize)
+ private int packedBlockSize = -1;
+ private long[] packedBits = null;
+ private int packedBlockCount = 0;
+
/**
* Initialize ASHVectors with an array of ASH-compressed vectors.
*
@@ -67,6 +79,18 @@ public ASHVectors(AsymmetricHashing ash,
this.ash = ash;
this.compressedVectors = compressedVectors;
this.scorer = new ASHScorer(ash);
+
+ this.d = ash.quantizedDim;
+ this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
+
+ int n = compressedVectors.length;
+ this.residualNorms = new float[n];
+ this.dotXMu = new float[n];
+ for (int i = 0; i < n; i++) {
+ var v = compressedVectors[i];
+ residualNorms[i] = v.residualNorm;
+ dotXMu[i] = v.dotWithLandmark;
+ }
}
@Override
@@ -126,6 +150,386 @@ public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(
return node -> f.similarityTo(compressedVectors[node]);
}
+// ============================================================================
+// ASH block-oriented scoring
+// ============================================================================
+
+ /**
+ * Scalar reference block scorer.
+ *
+ *
+ * This scorer computes ASH scores by invoking the scalar per-vector
+ * score function for each vector in the requested range.
+ *
+ *
+ *
+ * It serves as a correctness reference and baseline implementation.
+ * No block-level optimizations are applied.
+ *
+ */
+ private static final class ScalarASHBlockScorer implements ASHBlockScorer {
+
+ private final ASHScorer.ASHScoreFunction scorer;
+ private final AsymmetricHashing.QuantizedVector[] vectors;
+
+ ScalarASHBlockScorer(
+ ASHScorer.ASHScoreFunction scorer,
+ AsymmetricHashing.QuantizedVector[] vectors) {
+ this.scorer = scorer;
+ this.vectors = vectors;
+ }
+
+ @Override
+ public void scoreRange(int start, int count, float[] out) {
+ if (count < 0) throw new IllegalArgumentException("count must be >= 0");
+ if (start < 0 || start + count > vectors.length) {
+ throw new IllegalArgumentException(
+ "Range out of bounds: [" + start + ", " + (start + count) + ")"
+ );
+ }
+ if (out.length < count) {
+ throw new IllegalArgumentException("out.length < count");
+ }
+ if (count == 0) return;
+
+ final int end = start + count;
+ for (int i = start, o = 0; i < end; i++, o++) {
+ out[o] = scorer.similarityTo(vectors[i]);
+ }
+ }
+ }
+
+ /**
+ * Returns a scalar reference block scorer.
+ *
+ *
+ * The returned scorer matches {@link #scoreFunctionFor} exactly
+ * and computes each score independently.
+ *
+ */
+ public ASHBlockScorer blockScorerFor(
+ VectorFloat> query,
+ VectorSimilarityFunction similarityFunction) {
+
+ if (similarityFunction != VectorSimilarityFunction.DOT_PRODUCT) {
+ throw new UnsupportedOperationException(
+ "ASH block scorer supports DOT_PRODUCT only");
+ }
+
+ ASHScorer.ASHScoreFunction f =
+ scorer.scoreFunctionFor(query, similarityFunction);
+
+ return new ScalarASHBlockScorer(f, compressedVectors);
+ }
+
+ /**
+ * Block-structured scalar ASH scorer.
+ *
+ *
+ * This scorer reorganizes the computation to operate on contiguous
+ * blocks of vectors while remaining fully scalar.
+ *
+ *
+ *
+ * The loop structure is:
+ *
+ *
+ * - outer loop over blocks of vectors
+ * - outer loop over binary words
+ * - inner loop over vectors within the block
+ *
+ *
+ *
+ * This improves cache locality and provides a stable structural
+ * foundation for block-oriented evaluation.
+ *
+ */
+ public ASHBlockScorer blockScorerFor(
+ VectorFloat> query,
+ VectorSimilarityFunction similarityFunction,
+ int blockSize) {
+
+ if (similarityFunction != VectorSimilarityFunction.DOT_PRODUCT) {
+ throw new UnsupportedOperationException(
+ "ASH block scorer supports DOT_PRODUCT only");
+ }
+ if (blockSize <= 0) {
+ throw new IllegalArgumentException("blockSize must be > 0");
+ }
+
+ return new BlockedScalarASHBlockScorer(ash, compressedVectors, query, blockSize);
+ }
+
+ /**
+ * Block-oriented ASH scorer using register-local accumulation.
+ *
+ *
+ * This implementation evaluates ASH scores for a contiguous block
+ * of vectors while accumulating intermediate values directly into
+ * local scalar variables.
+ *
+ *
+ *
+ * Compared to array-based accumulation, this approach:
+ *
+ *
+ * - reduces memory traffic
+ * - keeps hot values in registers
+ * - preserves exact ASH scoring semantics
+ *
+ *
+ *
+ * This scorer is suitable for brute-force scans and graph neighborhood
+ * evaluation where vectors are processed contiguously.
+ *
+ */
+ public ASHBlockScorer blockScorerRegisterAccFor(
+ VectorFloat> query,
+ VectorSimilarityFunction similarityFunction,
+ int blockSize) {
+
+ if (similarityFunction != VectorSimilarityFunction.DOT_PRODUCT) {
+ throw new UnsupportedOperationException("DOT_PRODUCT only");
+ }
+ if (blockSize <= 0) {
+ throw new IllegalArgumentException("blockSize must be > 0");
+ }
+
+ return new RegisterAccASHBlockScorer(
+ ash,
+ compressedVectors,
+ residualNorms,
+ dotXMu,
+ query
+ );
+ }
+
+ /**
+ * Block-structured scalar ASH scorer.
+ *
+ *
+ * This scorer performs scalar bit-walk accumulation for each vector,
+ * but organizes the computation by blocks and binary words.
+ *
+ *
+ *
+ * It preserves exact correctness and serves as a clear baseline
+ * for block-based implementations.
+ *
+ */
+ private static final class BlockedScalarASHBlockScorer implements ASHBlockScorer {
+
+ private final AsymmetricHashing ash;
+ private final AsymmetricHashing.QuantizedVector[] vectors;
+ private final int blockSize;
+
+ private final float[] tildeQ;
+ private final float sumTildeQ;
+ private final float dotQMu;
+ private final float muNormSq;
+ private final float invSqrtD;
+
+ private final int d;
+ private final int words;
+ private final int originalDim;
+
+ BlockedScalarASHBlockScorer(
+ AsymmetricHashing ash,
+ AsymmetricHashing.QuantizedVector[] vectors,
+ VectorFloat> query,
+ int blockSize) {
+
+ this.ash = ash;
+ this.vectors = vectors;
+ this.blockSize = blockSize;
+
+ this.d = ash.quantizedDim;
+ this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
+ this.originalDim = ash.originalDimension;
+
+ final VectorFloat> mu = ash.globalMean;
+
+ final float[] qArr = new float[originalDim];
+ final float[] muArr = new float[originalDim];
+ for (int i = 0; i < originalDim; i++) {
+ qArr[i] = query.get(i);
+ muArr[i] = mu.get(i);
+ }
+
+ this.tildeQ = new float[d];
+ final float[][] A = ash.stiefelTransform.AFloat;
+ for (int r = 0; r < d; r++) {
+ float acc = 0f;
+ float[] Arow = A[r];
+ for (int j = 0; j < originalDim; j++) {
+ acc += Arow[j] * (qArr[j] - muArr[j]);
+ }
+ tildeQ[r] = acc;
+ }
+
+ float s = 0f;
+ for (int i = 0; i < d; i++) s += tildeQ[i];
+ this.sumTildeQ = s;
+
+ this.dotQMu = VectorUtil.dotProduct(query, mu);
+ this.muNormSq = VectorUtil.dotProduct(mu, mu);
+ this.invSqrtD = (float) (1.0 / Math.sqrt(d));
+ }
+
+ @Override
+ public void scoreRange(int start, int count, float[] out) {
+ if (count < 0) throw new IllegalArgumentException("count must be >= 0");
+ if (start < 0 || start + count > vectors.length) {
+ throw new IllegalArgumentException(
+ "Range out of bounds: [" + start + ", " + (start + count) + ")"
+ );
+ }
+ if (out.length < count) {
+ throw new IllegalArgumentException("out.length < count");
+ }
+ if (count == 0) return;
+
+ final float[] acc = new float[Math.min(blockSize, count)];
+
+ int baseOrd = start;
+ int remaining = count;
+
+ while (remaining > 0) {
+ final int blockLen = Math.min(blockSize, remaining);
+
+ for (int i = 0; i < blockLen; i++) acc[i] = 0f;
+
+ for (int w = 0; w < words; w++) {
+ final int baseDim = w * 64;
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ long word = vectors[baseOrd + lane].binaryVector[w];
+
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+ if (idx < d) acc[lane] += tildeQ[idx];
+ word &= (word - 1);
+ }
+ }
+ }
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ final AsymmetricHashing.QuantizedVector v = vectors[baseOrd + lane];
+ final float scale = invSqrtD * v.residualNorm;
+
+ out[(baseOrd - start) + lane] =
+ scale * (2f * acc[lane] - sumTildeQ)
+ + dotQMu
+ + (v.dotWithLandmark - muNormSq);
+ }
+
+ baseOrd += blockLen;
+ remaining -= blockLen;
+ }
+ }
+ }
+
+ /**
+ * Block-oriented ASH scorer using register-local accumulation.
+ *
+ *
+ * This is the preferred block implementation. It avoids temporary
+ * per-block arrays and accumulates the masked-add term directly
+ * into scalar registers.
+ *
+ */
+ private static final class RegisterAccASHBlockScorer implements ASHBlockScorer {
+
+ private final AsymmetricHashing ash;
+ private final AsymmetricHashing.QuantizedVector[] vectors;
+ private final float[] residualNorms;
+ private final float[] dotXMu;
+
+ private final float[] tildeQ;
+ private final float sumTildeQ;
+ private final float dotQMu;
+ private final float muNormSq;
+ private final float invSqrtD;
+
+ private final int d;
+ private final int words;
+
+ RegisterAccASHBlockScorer(
+ AsymmetricHashing ash,
+ AsymmetricHashing.QuantizedVector[] vectors,
+ float[] residualNorms,
+ float[] dotXMu,
+ VectorFloat> query) {
+
+ this.ash = ash;
+ this.vectors = vectors;
+ this.residualNorms = residualNorms;
+ this.dotXMu = dotXMu;
+
+ this.d = ash.quantizedDim;
+ this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
+
+ final int D = ash.originalDimension;
+ final VectorFloat> mu = ash.globalMean;
+
+ final float[] qArr = new float[D];
+ final float[] muArr = new float[D];
+ for (int i = 0; i < D; i++) {
+ qArr[i] = query.get(i);
+ muArr[i] = mu.get(i);
+ }
+
+ this.tildeQ = new float[d];
+ final float[][] A = ash.stiefelTransform.AFloat;
+ for (int r = 0; r < d; r++) {
+ float s = 0f;
+ float[] Arow = A[r];
+ for (int j = 0; j < D; j++) {
+ s += Arow[j] * (qArr[j] - muArr[j]);
+ }
+ tildeQ[r] = s;
+ }
+
+ float sum = 0f;
+ for (int i = 0; i < d; i++) sum += tildeQ[i];
+ this.sumTildeQ = sum;
+
+ this.dotQMu = VectorUtil.dotProduct(query, mu);
+ this.muNormSq = VectorUtil.dotProduct(mu, mu);
+ this.invSqrtD = (float) (1.0 / Math.sqrt(d));
+ }
+
+ @Override
+ public void scoreRange(int start, int count, float[] out) {
+
+ for (int i = 0; i < count; i++) {
+ int idx = start + i;
+
+ float maskedAdd = 0f;
+ long[] bits = vectors[idx].binaryVector;
+
+ for (int w = 0; w < words; w++) {
+ long word = bits[w];
+ int baseDim = w * 64;
+
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int dim = baseDim + bit;
+ if (dim < d) maskedAdd += tildeQ[dim];
+ word &= (word - 1);
+ }
+ }
+
+ float scale = invSqrtD * residualNorms[idx];
+ out[i] =
+ scale * (2f * maskedAdd - sumTildeQ)
+ + dotQMu
+ + (dotXMu[idx] - muNormSq);
+ }
+ }
+ }
+
/**
* For ASH, precomputed and non-precomputed scoring are currently identical.
*
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 417d513b1..8d0ea7c45 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
@@ -256,26 +256,47 @@ default float ashDotRow(float[] Arow, float[] x) {
}
/**
- * maskedAdd = for b ∈ {0,1}^d stored as bitpacked longs.
- *
- * Contract: d <= tildeQ.length.
- * Default is scalar; SIMD backends may override.
+ * Computes maskedAdd for a block slice of vectors stored in packed block-column-major form.
+ * For each lane in [0, blockLen), writes:
+ * outMaskedAdd[lane] =
+ * packedBits layout:
+ * packedBits[blockWordBase + w*blockSize + laneIndex] is 64-bit word w for that lane.
+ * Default implementation is scalar bit-walk with a register accumulator per lane.
+ * SIMD backends may override.
*/
- default float ashMaskedAdd(float[] tildeQ, long[] bits, int d) {
- assert d <= tildeQ.length : "d must be <= tildeQ.length";
-
- float sum = 0.0f;
- int base = 0;
-
- for (int w = 0; w < bits.length && base < d; w++, base += 64) {
- long word = bits[w];
- while (word != 0L) {
- int bit = Long.numberOfTrailingZeros(word);
- int idx = base + bit;
- if (idx < d) sum += tildeQ[idx];
- word &= (word - 1);
+ default void ashMaskedAddBlockAllWords(
+ float[] tildeQ,
+ int d,
+ long[] packedBits,
+ int packedBase, // blockId * nWords * blockSize
+ int nWords,
+ int blockSize,
+ int laneStart,
+ int blockLen,
+ float[] outMaskedAdd
+ ) {
+ assert laneStart >= 0;
+ assert blockLen >= 0;
+ assert laneStart + blockLen <= blockSize;
+ assert outMaskedAdd.length >= blockLen : "outMaskedAdd too small";
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ float maskedAdd = 0f; // register accumulator (critical!)
+ int laneIndex = laneStart + lane;
+
+ for (int w = 0; w < nWords; w++) {
+ long word = packedBits[packedBase + (w * blockSize) + laneIndex];
+ int baseDim = w * 64;
+
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+ if (idx < d) maskedAdd += tildeQ[idx];
+ word &= (word - 1);
+ }
}
+
+ outMaskedAdd[lane] = maskedAdd;
}
- return sum;
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index b56368cfe..874e2b54e 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -19,10 +19,13 @@
import io.github.jbellis.jvector.example.util.SiftLoader;
import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.quantization.ASHBlockScorer;
import io.github.jbellis.jvector.quantization.CompressedVectors;
-import io.github.jbellis.jvector.quantization.AsymmetricHashing;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.quantization.ASHVectors;
+
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
import io.github.jbellis.jvector.vector.VectorUtil;
import io.github.jbellis.jvector.vector.types.VectorFloat;
import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
@@ -34,12 +37,12 @@
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;
-
-
import static java.lang.Math.abs;
public class DistancesASH {
public static void testASHEncodings(String filenameBase, String filenameQueries) throws IOException {
+ // Block sizes to benchmark
+ final int[] BLOCK_SIZES = {16, 32, 64}; // 16, 32, and/or 64
List> vectors = SiftLoader.readFvecs(filenameBase);
List> queries = SiftLoader.readFvecs(filenameQueries);
@@ -73,7 +76,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// - Base vectors x are NOT globally normalized.
// - For encoding, we compute μ on raw x.
// - The residual (x − μ) is normalized ONCE inside the binarizer,
- // producing \hat{x} = (x − μ) / ||x − μ|| (Eq. 6).
+ // producing \hat{x} = (x − μ) / ||x − μ|| (ASH Paper, Eq. 6).
// - Queries may be L2-normalized in the benchmark (standard practice),
// but are NOT normalized inside the encoder.
// - No other normalization steps are applied.
@@ -110,6 +113,30 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
long endTime = System.nanoTime();
System.out.println("\tEncoding took " + (endTime - startTime) / 1e9 + " seconds");
+ ASHVectors ashVectors = (ASHVectors) ashVecs;
+ VectorFloat> q0 = finalQueries.get(0);
+
+ ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
+ ASHBlockScorer blk = ashVectors.blockScorerRegisterAccFor(q0, VectorSimilarityFunction.DOT_PRODUCT, 64);
+
+ float[] a = new float[128];
+ float[] b = new float[128];
+ ref.scoreRange(0, 128, a);
+ blk.scoreRange(0, 128, b);
+
+ final float EPS = 1e-6f;
+
+ for (int i = 0; i < 128; i++) {
+ float diff = Math.abs(a[i] - b[i]);
+ if (diff > EPS) {
+ throw new AssertionError(
+ "Mismatch at " + i +
+ ": ref=" + a[i] +
+ ", blk=" + b[i] +
+ ", diff=" + diff);
+ }
+ }
+
// ------------------------------------------------------------------
// Shared parallelization setup (executor-honoring)
// ------------------------------------------------------------------
@@ -162,7 +189,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
System.out.println("\tAverage absolute dot-product error = " + distanceError);
// ==================================================================
- // [2] ASH distance timing (Eq. 11 only)
+ // [2] Scalar ASH distance timing (reference)
// ==================================================================
List> ashTasks = new ArrayList<>();
@@ -195,7 +222,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
long ashEnd = System.nanoTime();
- System.out.println("\tASH dot-product computations took "
+ System.out.println("\tScalar ASH dot-product computations took "
+ (ashEnd - ashStart) / 1e9 + " seconds");
// ==================================================================
@@ -235,6 +262,130 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Prevent dead-code elimination
System.out.println("\tdummyAccumulator = " + (float) (ashDummy + floatDummy));
System.out.println("--");
+ // ==================================================================
+ // [4] Block scoring timing: block baseline vs register-accumulator
+ // ==================================================================
+ for (int blockSize : BLOCK_SIZES) {
+
+ // --------------------------------------------------------------
+ // (A) Block baseline (scalar, block-structured)
+ // --------------------------------------------------------------
+ List> baselineTasks = new ArrayList<>();
+ long baselineStart = System.nanoTime();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ baselineTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
+ final float[] scores = new float[blockSize];
+
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+
+ ASHBlockScorer blockScorer =
+ ashVectors.blockScorerFor(
+ q,
+ VectorSimilarityFunction.DOT_PRODUCT,
+ blockSize);
+
+ int j = 0;
+ while (j + blockSize <= nVectors) {
+ blockScorer.scoreRange(j, blockSize, scores);
+ for (int k = 0; k < blockSize; k++) {
+ localSum += scores[k];
+ }
+ j += blockSize;
+ }
+
+ // Tail (scalar fallback)
+ if (j < nVectors) {
+ ScoreFunction.ApproximateScoreFunction f =
+ ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ for (; j < nVectors; j++) {
+ localSum += f.similarityTo(j);
+ }
+ }
+ }
+ return localSum;
+ }));
+ }
+
+ double baselineDummy = 0.0;
+ for (ForkJoinTask t : baselineTasks) {
+ baselineDummy += t.join();
+ }
+
+ long baselineEnd = System.nanoTime();
+
+ System.out.println(
+ "\tBlock baseline (blockSize=" + blockSize + ") took "
+ + (baselineEnd - baselineStart) / 1e9 + " seconds");
+
+ // --------------------------------------------------------------
+ // (B) Block register-accumulator scorer (fast block path)
+ // --------------------------------------------------------------
+ List> regAccTasks = new ArrayList<>();
+ long regAccStart = System.nanoTime();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ regAccTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
+ final float[] scores = new float[blockSize];
+
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+
+ ASHBlockScorer blockScorer =
+ ashVectors.blockScorerRegisterAccFor(
+ q,
+ VectorSimilarityFunction.DOT_PRODUCT,
+ blockSize);
+
+ int j = 0;
+ while (j + blockSize <= nVectors) {
+ blockScorer.scoreRange(j, blockSize, scores);
+ for (int k = 0; k < blockSize; k++) {
+ localSum += scores[k];
+ }
+ j += blockSize;
+ }
+
+ // Tail (scalar fallback)
+ if (j < nVectors) {
+ ScoreFunction.ApproximateScoreFunction f =
+ ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ for (; j < nVectors; j++) {
+ localSum += f.similarityTo(j);
+ }
+ }
+ }
+ return localSum;
+ }));
+ }
+
+ double regAccDummy = 0.0;
+ for (ForkJoinTask t : regAccTasks) {
+ regAccDummy += t.join();
+ }
+
+ long regAccEnd = System.nanoTime();
+
+ System.out.println(
+ "\tBlock register-acc (blockSize=" + blockSize + ") took "
+ + (regAccEnd - regAccStart) / 1e9 + " seconds");
+
+ // Prevent dead-code elimination across variants
+ ashDummy += baselineDummy + regAccDummy;
+ }
+
+ // Prevent dead-code elimination across all sections
+ System.out.println("\tdummyAccumulator = " + (float) (ashDummy + floatDummy));
+ System.out.println("--");
}
public static void runSIFT() throws IOException {
@@ -289,7 +440,7 @@ public static void runOpenai3072() throws IOException {
System.out.println("Running text-embedding-3-large_3072");
var baseVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
- var queryVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
+ var queryVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_query_vectors_10000.fvec";
testASHEncodings(baseVectors, queryVectors);
}
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 ea2211e8f..80808252a 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
@@ -30,6 +30,8 @@
import java.util.List;
class PanamaVectorUtilSupport implements VectorUtilSupport {
+ private static final String ASH_BLOCK_KERNEL =
+ System.getProperty("jvector.ash.blockKernel", "mask"); // "lut" or "mask"
static final int PREFERRED_BIT_SIZE = FloatVector.SPECIES_PREFERRED.vectorBitSize();
static final IntVector BYTE_TO_INT_MASK_512 = IntVector.broadcast(IntVector.SPECIES_512, 0xff);
@@ -1570,4 +1572,273 @@ public float ashDotRow(float[] Arow, float[] x) {
return acc;
}
+
+ @Override
+ public void ashMaskedAddBlockAllWords(float[] tildeQ,
+ int d,
+ long[] packedBits,
+ int blockWordBase,
+ int words,
+ int blockSize,
+ int laneStart,
+ int blockLen,
+ float[] outMaskedAdd) {
+
+ // Like in the ASH paper, SIMD lanes are dimensions (SPECIES_PREFERRED),
+ // masked load from tildeQ and horizontal sum per neighbor.
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int LANES = SPEC.length();
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ float maskedAdd = 0f; // register scalar accumulator
+ int laneIndex = laneStart + lane;
+
+ for (int w = 0; w < words; w++) {
+ long wordBits = packedBits[blockWordBase + (w * blockSize) + laneIndex];
+ int baseDim = w * 64;
+
+ // process this 64-bit word in LANES-bit chunks
+ for (int off = 0; off < 64 && (baseDim + off) < d; off += LANES) {
+ int qBase = baseDim + off;
+
+ // Range mask for tail dims beyond d
+ VectorMask inRange = SPEC.indexInRange(qBase, d);
+
+ long maskBits;
+ if (LANES == 16) {
+ maskBits = (wordBits >>> off) & 0xFFFFL;
+ } else if (LANES == 8) {
+ maskBits = (wordBits >>> off) & 0xFFL;
+ } else if (LANES == 4) {
+ maskBits = (wordBits >>> off) & 0xFL;
+ } else {
+ // Unusual width; fall back to scalar for this chunk
+ long m = wordBits >>> off;
+ float s = 0f;
+ for (int b = 0; b < LANES; b++) {
+ if (((m >>> b) & 1L) != 0L) {
+ int idx = qBase + b;
+ if (idx < d) s += tildeQ[idx];
+ }
+ }
+ maskedAdd += s;
+ continue;
+ }
+
+ if (maskBits == 0L) continue;
+
+ VectorMask bitMask =
+ VectorMask.fromLong(SPEC, maskBits).and(inRange);
+
+ // masked load (paper)
+ FloatVector masked = FloatVector.fromArray(SPEC, tildeQ, qBase, bitMask);
+
+ // horizontal sum (paper)
+ maskedAdd += masked.reduceLanes(VectorOperators.ADD);
+ }
+ }
+
+ outMaskedAdd[lane] = maskedAdd;
+ }
+ }
+
+ private static void ashMaskedAddBlockWordLutNibble(float[] tildeQ,
+ int baseDim,
+ int d,
+ long[] packedBits,
+ int packedBase,
+ int laneStart,
+ int blockLen,
+ float[] acc) {
+ for (int off = 0; off < 64 && (baseDim + off) < d; off += 4) {
+ int q0 = baseDim + off;
+
+ float v0 = (q0 + 0 < d) ? tildeQ[q0 + 0] : 0f;
+ float v1 = (q0 + 1 < d) ? tildeQ[q0 + 1] : 0f;
+ float v2 = (q0 + 2 < d) ? tildeQ[q0 + 2] : 0f;
+ float v3 = (q0 + 3 < d) ? tildeQ[q0 + 3] : 0f;
+
+ float lut1 = v0;
+ float lut2 = v1;
+ float lut3 = v0 + v1;
+ float lut4 = v2;
+ float lut5 = v0 + v2;
+ float lut6 = v1 + v2;
+ float lut7 = v0 + v1 + v2;
+ float lut8 = v3;
+ float lut9 = v0 + v3;
+ float lut10 = v1 + v3;
+ float lut11 = v0 + v1 + v3;
+ float lut12 = v2 + v3;
+ float lut13 = v0 + v2 + v3;
+ float lut14 = v1 + v2 + v3;
+ float lut15 = v0 + v1 + v2 + v3;
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ long word = packedBits[packedBase + laneStart + lane];
+ int mask = (int) ((word >>> off) & 0xF);
+ switch (mask) {
+ case 0: break;
+ case 1: acc[lane] += lut1; break;
+ case 2: acc[lane] += lut2; break;
+ case 3: acc[lane] += lut3; break;
+ case 4: acc[lane] += lut4; break;
+ case 5: acc[lane] += lut5; break;
+ case 6: acc[lane] += lut6; break;
+ case 7: acc[lane] += lut7; break;
+ case 8: acc[lane] += lut8; break;
+ case 9: acc[lane] += lut9; break;
+ case 10: acc[lane] += lut10; break;
+ case 11: acc[lane] += lut11; break;
+ case 12: acc[lane] += lut12; break;
+ case 13: acc[lane] += lut13; break;
+ case 14: acc[lane] += lut14; break;
+ case 15: acc[lane] += lut15; break;
+ }
+ }
+ }
+ }
+
+ private static void ashMaskedAddBlockWordLutByteTwoNibbles(float[] tildeQ,
+ int baseDim,
+ int d,
+ long[] packedBits,
+ int packedBase,
+ int laneStart,
+ int blockLen,
+ float[] acc) {
+ for (int off = 0; off < 64 && (baseDim + off) < d; off += 8) {
+ int q0 = baseDim + off;
+
+ float v0 = (q0 + 0 < d) ? tildeQ[q0 + 0] : 0f;
+ float v1 = (q0 + 1 < d) ? tildeQ[q0 + 1] : 0f;
+ float v2 = (q0 + 2 < d) ? tildeQ[q0 + 2] : 0f;
+ float v3 = (q0 + 3 < d) ? tildeQ[q0 + 3] : 0f;
+ float v4 = (q0 + 4 < d) ? tildeQ[q0 + 4] : 0f;
+ float v5 = (q0 + 5 < d) ? tildeQ[q0 + 5] : 0f;
+ float v6 = (q0 + 6 < d) ? tildeQ[q0 + 6] : 0f;
+ float v7 = (q0 + 7 < d) ? tildeQ[q0 + 7] : 0f;
+
+ // Low nibble (v0..v3)
+ float lo1 = v0;
+ float lo2 = v1;
+ float lo3 = v0 + v1;
+ float lo4 = v2;
+ float lo5 = v0 + v2;
+ float lo6 = v1 + v2;
+ float lo7 = v0 + v1 + v2;
+ float lo8 = v3;
+ float lo9 = v0 + v3;
+ float lo10 = v1 + v3;
+ float lo11 = v0 + v1 + v3;
+ float lo12 = v2 + v3;
+ float lo13 = v0 + v2 + v3;
+ float lo14 = v1 + v2 + v3;
+ float lo15 = v0 + v1 + v2 + v3;
+
+ // High nibble (v4..v7)
+ float hi1 = v4;
+ float hi2 = v5;
+ float hi3 = v4 + v5;
+ float hi4 = v6;
+ float hi5 = v4 + v6;
+ float hi6 = v5 + v6;
+ float hi7 = v4 + v5 + v6;
+ float hi8 = v7;
+ float hi9 = v4 + v7;
+ float hi10 = v5 + v7;
+ float hi11 = v4 + v5 + v7;
+ float hi12 = v6 + v7;
+ float hi13 = v4 + v6 + v7;
+ float hi14 = v5 + v6 + v7;
+ float hi15 = v4 + v5 + v6 + v7;
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ long word = packedBits[packedBase + laneStart + lane];
+ int mask = (int) ((word >>> off) & 0xFF);
+
+ int lo = mask & 0xF;
+ int hi = (mask >>> 4) & 0xF;
+
+ float addLo = 0f;
+ switch (lo) {
+ case 0: break;
+ case 1: addLo = lo1; break;
+ case 2: addLo = lo2; break;
+ case 3: addLo = lo3; break;
+ case 4: addLo = lo4; break;
+ case 5: addLo = lo5; break;
+ case 6: addLo = lo6; break;
+ case 7: addLo = lo7; break;
+ case 8: addLo = lo8; break;
+ case 9: addLo = lo9; break;
+ case 10: addLo = lo10; break;
+ case 11: addLo = lo11; break;
+ case 12: addLo = lo12; break;
+ case 13: addLo = lo13; break;
+ case 14: addLo = lo14; break;
+ case 15: addLo = lo15; break;
+ }
+
+ float addHi = 0f;
+ switch (hi) {
+ case 0: break;
+ case 1: addHi = hi1; break;
+ case 2: addHi = hi2; break;
+ case 3: addHi = hi3; break;
+ case 4: addHi = hi4; break;
+ case 5: addHi = hi5; break;
+ case 6: addHi = hi6; break;
+ case 7: addHi = hi7; break;
+ case 8: addHi = hi8; break;
+ case 9: addHi = hi9; break;
+ case 10: addHi = hi10; break;
+ case 11: addHi = hi11; break;
+ case 12: addHi = hi12; break;
+ case 13: addHi = hi13; break;
+ case 14: addHi = hi14; break;
+ case 15: addHi = hi15; break;
+ }
+
+ acc[lane] += (addLo + addHi);
+ }
+ }
+ }
+
+ private static void ashMaskedAddBlockWordMask(float[] tildeQ,
+ int baseDim,
+ int d,
+ long[] packedBits,
+ int packedBase,
+ int laneStart,
+ int blockLen,
+ float[] acc) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int lanes = SPEC.length();
+
+ for (int off = 0; off < 64 && (baseDim + off) < d; off += lanes) {
+ int qBase = baseDim + off;
+
+ VectorMask inRange = SPEC.indexInRange(qBase, d);
+ FloatVector qVec = FloatVector.fromArray(SPEC, tildeQ, qBase, inRange);
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ long word = packedBits[packedBase + laneStart + lane];
+
+ long maskBits;
+ if (lanes == 16) {
+ maskBits = (word >>> off) & 0xFFFFL;
+ } else if (lanes == 8) {
+ maskBits = (word >>> off) & 0xFFL;
+ } else { // NEON 4 lanes
+ maskBits = (word >>> off) & 0xFL;
+ }
+
+ if (maskBits == 0L) continue;
+
+ VectorMask bitMask = VectorMask.fromLong(SPEC, maskBits).and(inRange);
+ acc[lane] += qVec.reduceLanes(VectorOperators.ADD, bitMask);
+ }
+ }
+ }
}
From 78f67acd09bcb9b9382d58690369ddd5670577cf Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Tue, 30 Dec 2025 23:23:58 +0000
Subject: [PATCH 007/243] Improved backend selection logic and SIMD
performance. Also improved DistancesASH.
---
.../jvector/quantization/ASHVectors.java | 339 ++++++++++--------
.../jvector/vector/VectorUtilSupport.java | 7 +
.../jbellis/jvector/example/DistancesASH.java | 148 +++++---
.../vector/PanamaVectorUtilSupport.java | 5 +
4 files changed, 314 insertions(+), 185 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index 5ee690c6e..86277d4d2 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -21,6 +21,8 @@
import io.github.jbellis.jvector.util.RamUsageEstimator;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.VectorUtilSupport;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
import io.github.jbellis.jvector.vector.types.VectorFloat;
import java.io.DataOutput;
@@ -55,6 +57,19 @@ public class ASHVectors implements CompressedVectors {
final AsymmetricHashing.QuantizedVector[] compressedVectors;
final ASHScorer scorer;
+ enum AshBlockKernel {
+ AUTO,
+ SCALAR,
+ SIMD;
+
+ static AshBlockKernel fromProperty() {
+ String v = System.getProperty("jvector.ash.blockKernel", "auto").toLowerCase();
+ if ("scalar".equals(v)) return SCALAR;
+ if ("simd".equals(v)) return SIMD;
+ return AUTO;
+ }
+ }
+
// Cached scalar headers (avoid object chasing in blocked scorer)
private final float[] residualNorms;
private final float[] dotXMu;
@@ -223,67 +238,33 @@ public ASHBlockScorer blockScorerFor(
}
/**
- * Block-structured scalar ASH scorer.
+ * Returns a block-oriented ASH scorer for contiguous vectors.
*
*
- * This scorer reorganizes the computation to operate on contiguous
- * blocks of vectors while remaining fully scalar.
+ * This factory selects between a portable scalar implementation and a
+ * SIMD implementation of ASH block scoring (ASH paper, Figure 1 / Table 1).
*
*
*
- * The loop structure is:
+ * Selection is controlled by the JVM property:
*
- *
- * - outer loop over blocks of vectors
- * - outer loop over binary words
- * - inner loop over vectors within the block
- *
*
- *
- * This improves cache locality and provides a stable structural
- * foundation for block-oriented evaluation.
- *
- */
- public ASHBlockScorer blockScorerFor(
- VectorFloat> query,
- VectorSimilarityFunction similarityFunction,
- int blockSize) {
-
- if (similarityFunction != VectorSimilarityFunction.DOT_PRODUCT) {
- throw new UnsupportedOperationException(
- "ASH block scorer supports DOT_PRODUCT only");
- }
- if (blockSize <= 0) {
- throw new IllegalArgumentException("blockSize must be > 0");
- }
-
- return new BlockedScalarASHBlockScorer(ash, compressedVectors, query, blockSize);
- }
-
- /**
- * Block-oriented ASH scorer using register-local accumulation.
+ *
+ * -Djvector.ash.blockKernel=auto|scalar|simd
+ *
*
- *
- * This implementation evaluates ASH scores for a contiguous block
- * of vectors while accumulating intermediate values directly into
- * local scalar variables.
- *
- *
- *
- * Compared to array-based accumulation, this approach:
- *
*
- * - reduces memory traffic
- * - keeps hot values in registers
- * - preserves exact ASH scoring semantics
+ * - auto (default): the active VectorUtilSupport backend determines whether execution is vectorized or scalar
+ * - scalar: force the portable scalar register-accumulating implementation
+ * - simd: simd: force the block scorer defined in the ASH paper; execution may still be scalar if the active backend does not support masked-load SIMD
*
*
*
- * This scorer is suitable for brute-force scans and graph neighborhood
- * evaluation where vectors are processed contiguously.
+ * In all cases, accumulation of the masked-add term ⟨q̃, b⟩ is performed in
+ * local scalar registers, ensuring numerical equivalence across backends.
*
*/
- public ASHBlockScorer blockScorerRegisterAccFor(
+ public ASHBlockScorer blockScorerFor(
VectorFloat> query,
VectorSimilarityFunction similarityFunction,
int blockSize) {
@@ -295,33 +276,84 @@ public ASHBlockScorer blockScorerRegisterAccFor(
throw new IllegalArgumentException("blockSize must be > 0");
}
- return new RegisterAccASHBlockScorer(
+ final AshBlockKernel kernel = AshBlockKernel.fromProperty();
+
+ // ------------------------------------------------------------
+ // Forced scalar path (portable register-accumulator)
+ // ------------------------------------------------------------
+ if (kernel == AshBlockKernel.SCALAR) {
+ return new BlockScalarASHBlockScorer(
+ ash,
+ compressedVectors,
+ residualNorms,
+ dotXMu,
+ query
+ );
+ }
+
+ // ------------------------------------------------------------
+ // SIMD or AUTO path (paper-faithful block kernel)
+ // ------------------------------------------------------------
+ // AUTO and SIMD intentionally share the same implementation.
+ // VectorUtilSupport decides whether masked-load is vectorized
+ // or falls back to scalar internally.
+ ensurePackedBits(blockSize);
+
+ return new BlockSimdASHBlockScorer(
ash,
compressedVectors,
residualNorms,
dotXMu,
+ packedBits,
+ blockSize,
query
);
}
+ private void ensurePackedBits(int blockSize) {
+ if (blockSize <= 0) {
+ throw new IllegalArgumentException("blockSize must be > 0");
+ }
+ if (packedBits != null && packedBlockSize == blockSize) {
+ return;
+ }
+
+ final int n = compressedVectors.length;
+ packedBlockSize = blockSize;
+ packedBlockCount = (n + blockSize - 1) / blockSize;
+ packedBits = new long[packedBlockCount * words * blockSize];
+
+ for (int b = 0; b < packedBlockCount; b++) {
+ int baseOrd = b * blockSize;
+ int blockLen = Math.min(blockSize, n - baseOrd);
+
+ for (int w = 0; w < words; w++) {
+ int dstBase = (b * words + w) * blockSize;
+ for (int lane = 0; lane < blockLen; lane++) {
+ packedBits[dstBase + lane] =
+ compressedVectors[baseOrd + lane].binaryVector[w];
+ }
+ // remaining lanes default to 0
+ }
+ }
+ }
+
+
/**
- * Block-structured scalar ASH scorer.
- *
- *
- * This scorer performs scalar bit-walk accumulation for each vector,
- * but organizes the computation by blocks and binary words.
- *
+ * Block-oriented ASH scorer using register-local accumulation.
*
*
- * It preserves exact correctness and serves as a clear baseline
- * for block-based implementations.
+ * This is the preferred block implementation. It avoids temporary
+ * per-block arrays and accumulates the masked-add term directly
+ * into scalar registers.
*
*/
- private static final class BlockedScalarASHBlockScorer implements ASHBlockScorer {
+ private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
private final AsymmetricHashing ash;
private final AsymmetricHashing.QuantizedVector[] vectors;
- private final int blockSize;
+ private final float[] residualNorms;
+ private final float[] dotXMu;
private final float[] tildeQ;
private final float sumTildeQ;
@@ -331,27 +363,28 @@ private static final class BlockedScalarASHBlockScorer implements ASHBlockScorer
private final int d;
private final int words;
- private final int originalDim;
- BlockedScalarASHBlockScorer(
+ BlockScalarASHBlockScorer(
AsymmetricHashing ash,
AsymmetricHashing.QuantizedVector[] vectors,
- VectorFloat> query,
- int blockSize) {
+ float[] residualNorms,
+ float[] dotXMu,
+ VectorFloat> query) {
this.ash = ash;
this.vectors = vectors;
- this.blockSize = blockSize;
+ this.residualNorms = residualNorms;
+ this.dotXMu = dotXMu;
this.d = ash.quantizedDim;
this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
- this.originalDim = ash.originalDimension;
+ final int D = ash.originalDimension;
final VectorFloat> mu = ash.globalMean;
- final float[] qArr = new float[originalDim];
- final float[] muArr = new float[originalDim];
- for (int i = 0; i < originalDim; i++) {
+ final float[] qArr = new float[D];
+ final float[] muArr = new float[D];
+ for (int i = 0; i < D; i++) {
qArr[i] = query.get(i);
muArr[i] = mu.get(i);
}
@@ -359,17 +392,17 @@ private static final class BlockedScalarASHBlockScorer implements ASHBlockScorer
this.tildeQ = new float[d];
final float[][] A = ash.stiefelTransform.AFloat;
for (int r = 0; r < d; r++) {
- float acc = 0f;
+ float s = 0f;
float[] Arow = A[r];
- for (int j = 0; j < originalDim; j++) {
- acc += Arow[j] * (qArr[j] - muArr[j]);
+ for (int j = 0; j < D; j++) {
+ s += Arow[j] * (qArr[j] - muArr[j]);
}
- tildeQ[r] = acc;
+ tildeQ[r] = s;
}
- float s = 0f;
- for (int i = 0; i < d; i++) s += tildeQ[i];
- this.sumTildeQ = s;
+ float sum = 0f;
+ for (int i = 0; i < d; i++) sum += tildeQ[i];
+ this.sumTildeQ = sum;
this.dotQMu = VectorUtil.dotProduct(query, mu);
this.muNormSq = VectorUtil.dotProduct(mu, mu);
@@ -378,98 +411,107 @@ private static final class BlockedScalarASHBlockScorer implements ASHBlockScorer
@Override
public void scoreRange(int start, int count, float[] out) {
- if (count < 0) throw new IllegalArgumentException("count must be >= 0");
- if (start < 0 || start + count > vectors.length) {
- throw new IllegalArgumentException(
- "Range out of bounds: [" + start + ", " + (start + count) + ")"
- );
- }
- if (out.length < count) {
- throw new IllegalArgumentException("out.length < count");
- }
- if (count == 0) return;
-
- final float[] acc = new float[Math.min(blockSize, count)];
- int baseOrd = start;
- int remaining = count;
-
- while (remaining > 0) {
- final int blockLen = Math.min(blockSize, remaining);
+ for (int i = 0; i < count; i++) {
+ int idx = start + i;
- for (int i = 0; i < blockLen; i++) acc[i] = 0f;
+ float maskedAdd = 0f;
+ long[] bits = vectors[idx].binaryVector;
for (int w = 0; w < words; w++) {
- final int baseDim = w * 64;
-
- for (int lane = 0; lane < blockLen; lane++) {
- long word = vectors[baseOrd + lane].binaryVector[w];
+ long word = bits[w];
+ int baseDim = w * 64;
- while (word != 0L) {
- int bit = Long.numberOfTrailingZeros(word);
- int idx = baseDim + bit;
- if (idx < d) acc[lane] += tildeQ[idx];
- word &= (word - 1);
- }
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int dim = baseDim + bit;
+ if (dim < d) maskedAdd += tildeQ[dim];
+ word &= (word - 1);
}
}
- for (int lane = 0; lane < blockLen; lane++) {
- final AsymmetricHashing.QuantizedVector v = vectors[baseOrd + lane];
- final float scale = invSqrtD * v.residualNorm;
-
- out[(baseOrd - start) + lane] =
- scale * (2f * acc[lane] - sumTildeQ)
- + dotQMu
- + (v.dotWithLandmark - muNormSq);
- }
-
- baseOrd += blockLen;
- remaining -= blockLen;
+ float scale = invSqrtD * residualNorms[idx];
+ out[i] =
+ scale * (2f * maskedAdd - sumTildeQ)
+ + dotQMu
+ + (dotXMu[idx] - muNormSq);
}
}
}
/**
- * Block-oriented ASH scorer using register-local accumulation.
+ * Block-oriented ASH scorer using SIMD masked-loads with register-local
+ * accumulation.
*
*
- * This is the preferred block implementation. It avoids temporary
- * per-block arrays and accumulates the masked-add term directly
- * into scalar registers.
+ * This scorer is a faithful implementation of Figure 1
+ * and Table 1 kernel from the ASH paper
+ *
+ *
+ *
+ * - binary codes are stored in block column-major layout
+ * - SIMD lanes correspond to dimensions
+ * - masked loads select active dimensions
+ * - horizontal sums are accumulated per vector
+ *
+ *
+ *
+ * This implementation accumulates masked-add results directly into scalar registers,
+ * avoiding per-lane memory traffic and restoring performance.
*
*/
- private static final class RegisterAccASHBlockScorer implements ASHBlockScorer {
+ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private final AsymmetricHashing ash;
private final AsymmetricHashing.QuantizedVector[] vectors;
+
private final float[] residualNorms;
private final float[] dotXMu;
+ private final long[] packedBits;
+ private final int blockSize;
+
+ private final int d;
+ private final int words;
+
+ // Query precompute
private final float[] tildeQ;
private final float sumTildeQ;
private final float dotQMu;
private final float muNormSq;
private final float invSqrtD;
- private final int d;
- private final int words;
+ // Scratch (one per scorer instance)
+ private final float[] maskedAdds;
+
+ private final VectorUtilSupport vecUtil;
- RegisterAccASHBlockScorer(
+ BlockSimdASHBlockScorer(
AsymmetricHashing ash,
AsymmetricHashing.QuantizedVector[] vectors,
float[] residualNorms,
float[] dotXMu,
+ long[] packedBits,
+ int blockSize,
VectorFloat> query) {
+ this.vecUtil =
+ io.github.jbellis.jvector.vector.VectorizationProvider
+ .getInstance()
+ .getVectorUtilSupport();
this.ash = ash;
this.vectors = vectors;
this.residualNorms = residualNorms;
this.dotXMu = dotXMu;
+ this.packedBits = packedBits;
+ this.blockSize = blockSize;
this.d = ash.quantizedDim;
this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
+ this.maskedAdds = new float[blockSize];
+
+ // ---- Query precompute (identical to ASHScorer) ----
final int D = ash.originalDimension;
final VectorFloat> mu = ash.globalMean;
@@ -502,30 +544,45 @@ private static final class RegisterAccASHBlockScorer implements ASHBlockScorer {
@Override
public void scoreRange(int start, int count, float[] out) {
+ if (count == 0) return;
- for (int i = 0; i < count; i++) {
- int idx = start + i;
-
- float maskedAdd = 0f;
- long[] bits = vectors[idx].binaryVector;
+ int ord = start;
+ int end = start + count;
+
+ while (ord < end) {
+ final int blockId = ord / blockSize;
+ final int blockBase = blockId * blockSize;
+ final int laneStart = ord - blockBase;
+ final int blockLen = Math.min(blockSize - laneStart, end - ord);
+
+ final int blockWordBase = blockId * words * blockSize;
+
+ // SIMD masked-add with register accumulation
+ vecUtil.ashMaskedAddBlockAllWords(
+ tildeQ,
+ d,
+ packedBits,
+ blockWordBase,
+ words,
+ blockSize,
+ laneStart,
+ blockLen,
+ maskedAdds
+ );
- for (int w = 0; w < words; w++) {
- long word = bits[w];
- int baseDim = w * 64;
+ // Finish Eq. 11 per vector
+ for (int lane = 0; lane < blockLen; lane++) {
+ int idx = ord + lane;
+ float scale = invSqrtD * residualNorms[idx];
+ float m = maskedAdds[lane];
- while (word != 0L) {
- int bit = Long.numberOfTrailingZeros(word);
- int dim = baseDim + bit;
- if (dim < d) maskedAdd += tildeQ[dim];
- word &= (word - 1);
- }
+ out[idx - start] =
+ scale * (2f * m - sumTildeQ)
+ + dotQMu
+ + (dotXMu[idx] - muNormSq);
}
- float scale = invSqrtD * residualNorms[idx];
- out[i] =
- scale * (2f * maskedAdd - sumTildeQ)
- + dotQMu
- + (dotXMu[idx] - muNormSq);
+ ord += blockLen;
}
}
}
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 8d0ea7c45..e40bcd862 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
@@ -239,6 +239,13 @@ default float pqDecodedCosineSimilarity(ByteSequence> encoded, int encodedOffs
// ASH kernels (default scalar implementations)
// ------------------------------------------------------------------
+ /**
+ * @return true if this backend provides a SIMD implementation of the masked-load kernel as described in the ASH paper.
+ */
+ default boolean supportsAshMaskedLoad() {
+ return false;
+ }
+
/**
* Dot product between a matrix row and a dense vector.
* Used by ASH projection: y_i = .
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 874e2b54e..8ca95c0d7 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -40,9 +40,16 @@
import static java.lang.Math.abs;
public class DistancesASH {
+ private static void printScorerInfo(String label, ASHBlockScorer scorer) {
+ System.out.println(
+ "\t[" + label + "] scorer implementation = "
+ + scorer.getClass().getName()
+ );
+ }
+
public static void testASHEncodings(String filenameBase, String filenameQueries) throws IOException {
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {16, 32, 64}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {32}; // 16, 32, and/or 64
List> vectors = SiftLoader.readFvecs(filenameBase);
List> queries = SiftLoader.readFvecs(filenameQueries);
@@ -117,7 +124,11 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
VectorFloat> q0 = finalQueries.get(0);
ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
- ASHBlockScorer blk = ashVectors.blockScorerRegisterAccFor(q0, VectorSimilarityFunction.DOT_PRODUCT, 64);
+ printScorerInfo("scalar-ref", ref);
+
+ ASHBlockScorer blk = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT, 64);
+ printScorerInfo("block-reg-acc", blk);
+
float[] a = new float[128];
float[] b = new float[128];
@@ -262,37 +273,75 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Prevent dead-code elimination
System.out.println("\tdummyAccumulator = " + (float) (ashDummy + floatDummy));
System.out.println("--");
+
// ==================================================================
- // [4] Block scoring timing: block baseline vs register-accumulator
+ // [4] Block scoring timing: block scalar vs block SIMD
// ==================================================================
for (int blockSize : BLOCK_SIZES) {
// --------------------------------------------------------------
- // (A) Block baseline (scalar, block-structured)
+ // Print scorer implementations ONCE (diagnostic)
// --------------------------------------------------------------
- List> baselineTasks = new ArrayList<>();
- long baselineStart = System.nanoTime();
+ {
+ VectorFloat> q = finalQueries.get(0);
+
+ // Scalar block scorer
+ System.setProperty("jvector.ash.blockKernel", "scalar");
+ ASHBlockScorer scalar =
+ ashVectors.blockScorerFor(
+ q,
+ VectorSimilarityFunction.DOT_PRODUCT,
+ blockSize
+ );
+ printScorerInfo(
+ "block-scalar (blockSize=" + blockSize + ")",
+ scalar
+ );
+
+ // SIMD block scorer
+ System.setProperty("jvector.ash.blockKernel", "simd");
+ ASHBlockScorer simd =
+ ashVectors.blockScorerFor(
+ q,
+ VectorSimilarityFunction.DOT_PRODUCT,
+ blockSize
+ );
+ printScorerInfo(
+ "block-simd (blockSize=" + blockSize + ")",
+ simd
+ );
+
+ // Restore default
+ System.clearProperty("jvector.ash.blockKernel");
+ }
+
+ // --------------------------------------------------------------
+ // (A) Block scalar timing
+ // --------------------------------------------------------------
+ System.setProperty("jvector.ash.blockKernel", "scalar");
+
+ List> scalarTasks = new ArrayList<>();
+ long scalarStart = System.nanoTime();
for (int start = 0; start < nQueries; start += chunkSize) {
final int s = start;
final int e = Math.min(start + chunkSize, nQueries);
- baselineTasks.add(simdExecutor.submit(() -> {
+ scalarTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
final float[] scores = new float[blockSize];
- for (int i = s; i < e; i++) {
- VectorFloat> q = finalQueries.get(i);
-
- ASHBlockScorer blockScorer =
+ for (int qi = s; qi < e; qi++) {
+ ASHBlockScorer scorer =
ashVectors.blockScorerFor(
- q,
+ finalQueries.get(qi),
VectorSimilarityFunction.DOT_PRODUCT,
- blockSize);
+ blockSize
+ );
int j = 0;
while (j + blockSize <= nVectors) {
- blockScorer.scoreRange(j, blockSize, scores);
+ scorer.scoreRange(j, blockSize, scores);
for (int k = 0; k < blockSize; k++) {
localSum += scores[k];
}
@@ -302,7 +351,10 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Tail (scalar fallback)
if (j < nVectors) {
ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ ashVecs.scoreFunctionFor(
+ finalQueries.get(qi),
+ VectorSimilarityFunction.DOT_PRODUCT
+ );
for (; j < nVectors; j++) {
localSum += f.similarityTo(j);
}
@@ -312,43 +364,45 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
}));
}
- double baselineDummy = 0.0;
- for (ForkJoinTask t : baselineTasks) {
- baselineDummy += t.join();
+ double scalarDummy = 0.0;
+ for (ForkJoinTask t : scalarTasks) {
+ scalarDummy += t.join();
}
- long baselineEnd = System.nanoTime();
+ long scalarEnd = System.nanoTime();
System.out.println(
- "\tBlock baseline (blockSize=" + blockSize + ") took "
- + (baselineEnd - baselineStart) / 1e9 + " seconds");
+ "\tBlock scalar (blockSize=" + blockSize + ") took "
+ + (scalarEnd - scalarStart) / 1e9 + " seconds"
+ );
// --------------------------------------------------------------
- // (B) Block register-accumulator scorer (fast block path)
+ // (B) Block SIMD timing
// --------------------------------------------------------------
- List> regAccTasks = new ArrayList<>();
- long regAccStart = System.nanoTime();
+ System.setProperty("jvector.ash.blockKernel", "simd");
+
+ List> simdTasks = new ArrayList<>();
+ long simdStart = System.nanoTime();
for (int start = 0; start < nQueries; start += chunkSize) {
final int s = start;
final int e = Math.min(start + chunkSize, nQueries);
- regAccTasks.add(simdExecutor.submit(() -> {
+ simdTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
final float[] scores = new float[blockSize];
- for (int i = s; i < e; i++) {
- VectorFloat> q = finalQueries.get(i);
-
- ASHBlockScorer blockScorer =
- ashVectors.blockScorerRegisterAccFor(
- q,
+ for (int qi = s; qi < e; qi++) {
+ ASHBlockScorer scorer =
+ ashVectors.blockScorerFor(
+ finalQueries.get(qi),
VectorSimilarityFunction.DOT_PRODUCT,
- blockSize);
+ blockSize
+ );
int j = 0;
while (j + blockSize <= nVectors) {
- blockScorer.scoreRange(j, blockSize, scores);
+ scorer.scoreRange(j, blockSize, scores);
for (int k = 0; k < blockSize; k++) {
localSum += scores[k];
}
@@ -358,7 +412,10 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Tail (scalar fallback)
if (j < nVectors) {
ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ ashVecs.scoreFunctionFor(
+ finalQueries.get(qi),
+ VectorSimilarityFunction.DOT_PRODUCT
+ );
for (; j < nVectors; j++) {
localSum += f.similarityTo(j);
}
@@ -368,19 +425,22 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
}));
}
- double regAccDummy = 0.0;
- for (ForkJoinTask t : regAccTasks) {
- regAccDummy += t.join();
+ double simdDummy = 0.0;
+ for (ForkJoinTask t : simdTasks) {
+ simdDummy += t.join();
}
- long regAccEnd = System.nanoTime();
+ long simdEnd = System.nanoTime();
System.out.println(
- "\tBlock register-acc (blockSize=" + blockSize + ") took "
- + (regAccEnd - regAccStart) / 1e9 + " seconds");
+ "\tBlock SIMD (blockSize=" + blockSize + ") took "
+ + (simdEnd - simdStart) / 1e9 + " seconds"
+ );
+
+ System.clearProperty("jvector.ash.blockKernel");
- // Prevent dead-code elimination across variants
- ashDummy += baselineDummy + regAccDummy;
+ // Prevent dead-code elimination
+ ashDummy += scalarDummy + simdDummy;
}
// Prevent dead-code elimination across all sections
@@ -415,8 +475,8 @@ public static void runCohere100k() throws IOException {
public static void runADA() throws IOException {
System.out.println("Running ada_002");
- var baseVectors = "./fvec/wikipedia_squad/100k/ada_002_100000_base_vectors.fvec";
- var queryVectors = "./fvec/wikipedia_squad/100k/ada_002_100000_query_vectors_10000.fvec";
+ var baseVectors = "./fvec/ada-002/ada_002_100000_base_vectors.fvec";
+ var queryVectors = "./fvec/ada-002/ada_002_100000_query_vectors_10000.fvec";
testASHEncodings(baseVectors, queryVectors);
}
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 80808252a..fbd9483eb 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
@@ -1573,6 +1573,11 @@ public float ashDotRow(float[] Arow, float[] x) {
return acc;
}
+ @Override
+ public boolean supportsAshMaskedLoad() {
+ return true;
+ }
+
@Override
public void ashMaskedAddBlockAllWords(float[] tildeQ,
int d,
From 2058037417131062d7a63f30a69ab1c8caedc99a Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Tue, 30 Dec 2025 23:59:58 +0000
Subject: [PATCH 008/243] Using only one backend selection per benchmark run in
DistancesASH.
---
.../jbellis/jvector/example/DistancesASH.java | 138 ++++--------------
1 file changed, 26 insertions(+), 112 deletions(-)
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 8ca95c0d7..0fac759fe 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -49,7 +49,7 @@ private static void printScorerInfo(String label, ASHBlockScorer scorer) {
public static void testASHEncodings(String filenameBase, String filenameQueries) throws IOException {
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {32}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {16, 32, 64, 128}; // 16, 32, and/or 64
List> vectors = SiftLoader.readFvecs(filenameBase);
List> queries = SiftLoader.readFvecs(filenameQueries);
@@ -126,8 +126,10 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
printScorerInfo("scalar-ref", ref);
- ASHBlockScorer blk = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT, 64);
- printScorerInfo("block-reg-acc", blk);
+ final int sanityBlockSize = BLOCK_SIZES[0];
+ ASHBlockScorer blk = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT, sanityBlockSize);
+ printScorerInfo("block", blk);
+
float[] a = new float[128];
@@ -275,59 +277,35 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
System.out.println("--");
// ==================================================================
- // [4] Block scoring timing: block scalar vs block SIMD
+ // [4] Block scoring timing (single mode per JVM run)
+ // NOTE: Choose the block kernel using:
+ // -Djvector.ash.blockKernel=scalar|simd|auto
// ==================================================================
+ final String kernelMode =
+ System.getProperty("jvector.ash.blockKernel", "auto").toLowerCase();
+
for (int blockSize : BLOCK_SIZES) {
- // --------------------------------------------------------------
- // Print scorer implementations ONCE (diagnostic)
- // --------------------------------------------------------------
+ // Print scorer implementation once per blockSize (diagnostic)
{
- VectorFloat> q = finalQueries.get(0);
-
- // Scalar block scorer
- System.setProperty("jvector.ash.blockKernel", "scalar");
- ASHBlockScorer scalar =
- ashVectors.blockScorerFor(
- q,
- VectorSimilarityFunction.DOT_PRODUCT,
- blockSize
- );
- printScorerInfo(
- "block-scalar (blockSize=" + blockSize + ")",
- scalar
- );
-
- // SIMD block scorer
- System.setProperty("jvector.ash.blockKernel", "simd");
- ASHBlockScorer simd =
+ ASHBlockScorer scorer =
ashVectors.blockScorerFor(
- q,
+ finalQueries.get(0),
VectorSimilarityFunction.DOT_PRODUCT,
blockSize
);
- printScorerInfo(
- "block-simd (blockSize=" + blockSize + ")",
- simd
- );
- // Restore default
- System.clearProperty("jvector.ash.blockKernel");
+ printScorerInfo(kernelMode + " (blockSize=" + blockSize + ")", scorer);
}
- // --------------------------------------------------------------
- // (A) Block scalar timing
- // --------------------------------------------------------------
- System.setProperty("jvector.ash.blockKernel", "scalar");
-
- List> scalarTasks = new ArrayList<>();
- long scalarStart = System.nanoTime();
+ List> blockTasks = new ArrayList<>();
+ long blockStart = System.nanoTime();
for (int start = 0; start < nQueries; start += chunkSize) {
final int s = start;
final int e = Math.min(start + chunkSize, nQueries);
- scalarTasks.add(simdExecutor.submit(() -> {
+ blockTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
final float[] scores = new float[blockSize];
@@ -348,7 +326,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
j += blockSize;
}
- // Tail (scalar fallback)
+ // Tail uses scalar per-vector scorer (correctness-preserving)
if (j < nVectors) {
ScoreFunction.ApproximateScoreFunction f =
ashVecs.scoreFunctionFor(
@@ -360,89 +338,25 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
}
}
}
- return localSum;
- }));
- }
-
- double scalarDummy = 0.0;
- for (ForkJoinTask t : scalarTasks) {
- scalarDummy += t.join();
- }
-
- long scalarEnd = System.nanoTime();
-
- System.out.println(
- "\tBlock scalar (blockSize=" + blockSize + ") took "
- + (scalarEnd - scalarStart) / 1e9 + " seconds"
- );
-
- // --------------------------------------------------------------
- // (B) Block SIMD timing
- // --------------------------------------------------------------
- System.setProperty("jvector.ash.blockKernel", "simd");
-
- List> simdTasks = new ArrayList<>();
- long simdStart = System.nanoTime();
-
- for (int start = 0; start < nQueries; start += chunkSize) {
- final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
-
- simdTasks.add(simdExecutor.submit(() -> {
- double localSum = 0.0;
- final float[] scores = new float[blockSize];
-
- for (int qi = s; qi < e; qi++) {
- ASHBlockScorer scorer =
- ashVectors.blockScorerFor(
- finalQueries.get(qi),
- VectorSimilarityFunction.DOT_PRODUCT,
- blockSize
- );
-
- int j = 0;
- while (j + blockSize <= nVectors) {
- scorer.scoreRange(j, blockSize, scores);
- for (int k = 0; k < blockSize; k++) {
- localSum += scores[k];
- }
- j += blockSize;
- }
- // Tail (scalar fallback)
- if (j < nVectors) {
- ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(
- finalQueries.get(qi),
- VectorSimilarityFunction.DOT_PRODUCT
- );
- for (; j < nVectors; j++) {
- localSum += f.similarityTo(j);
- }
- }
- }
return localSum;
}));
}
- double simdDummy = 0.0;
- for (ForkJoinTask t : simdTasks) {
- simdDummy += t.join();
+ double blockDummy = 0.0;
+ for (ForkJoinTask t : blockTasks) {
+ blockDummy += t.join();
}
- long simdEnd = System.nanoTime();
+ long blockEnd = System.nanoTime();
System.out.println(
- "\tBlock SIMD (blockSize=" + blockSize + ") took "
- + (simdEnd - simdStart) / 1e9 + " seconds"
+ "\tBlock " + kernelMode + " (blockSize=" + blockSize + ") took "
+ + (blockEnd - blockStart) / 1e9 + " seconds"
);
- System.clearProperty("jvector.ash.blockKernel");
-
- // Prevent dead-code elimination
- ashDummy += scalarDummy + simdDummy;
+ ashDummy += blockDummy;
}
-
// Prevent dead-code elimination across all sections
System.out.println("\tdummyAccumulator = " + (float) (ashDummy + floatDummy));
System.out.println("--");
From d194319f36f3502fba814028276982a48381d9ff Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Wed, 31 Dec 2025 16:55:37 +0000
Subject: [PATCH 009/243] New distance calculation benchmark for PQ.
---
.../jbellis/jvector/example/DistancesPQ.java | 302 ++++++++++++++++++
1 file changed, 302 insertions(+)
create mode 100644 jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
new file mode 100644
index 000000000..e57727fa9
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -0,0 +1,302 @@
+/*
+ * 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;
+
+import io.github.jbellis.jvector.example.util.SiftLoader;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.quantization.CompressedVectors;
+import io.github.jbellis.jvector.quantization.PQVectors;
+import io.github.jbellis.jvector.quantization.ProductQuantization;
+import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ForkJoinTask;
+
+public class DistancesPQ {
+
+ /**
+ * PQ benchmark aligned with DistancesASH:
+ *
+ * - Same datasets
+ * - Same query normalization
+ * - Same brute-force scan pattern
+ * - Same parallel executor
+ * - Same compression budget (≈ d/32 bytes per vector)
+ *
+ * Fast PQ path = PQDecoder (precomputedScoreFunctionFor).
+ */
+ public static void testPQEncodings(String filenameBase, String filenameQueries) throws IOException {
+
+ List> vectors = SiftLoader.readFvecs(filenameBase);
+ List> queries = SiftLoader.readFvecs(filenameQueries);
+
+ // ------------------------------------------------------------
+ // Remove zero-norm vectors (match DistancesASH)
+ // ------------------------------------------------------------
+ int beforeBase = vectors.size();
+ List> filteredBase = new ArrayList<>(vectors.size());
+ for (VectorFloat> v : vectors) {
+ if (VectorUtil.dotProduct(v, v) > 0f) {
+ filteredBase.add(v);
+ }
+ }
+ vectors = filteredBase;
+
+ int beforeQuery = queries.size();
+ List> filteredQueries = new ArrayList<>(queries.size());
+ for (VectorFloat> q : queries) {
+ if (VectorUtil.dotProduct(q, q) > 0f) {
+ filteredQueries.add(q);
+ }
+ }
+ queries = filteredQueries;
+
+ System.out.println("\tRemoved " + (beforeBase - vectors.size()) + " zero base vectors and "
+ + (beforeQuery - queries.size()) + " zero query vectors");
+
+ // Match ASH: normalize queries only
+ for (VectorFloat> q : queries) {
+ VectorUtil.l2normalize(q);
+ }
+
+ int dimension = vectors.get(0).length();
+
+ int nQueries = Math.min(1000, queries.size());
+ int nVectors = Math.min(100_000, vectors.size());
+
+ vectors = vectors.subList(0, nVectors);
+ queries = queries.subList(0, nQueries);
+
+ final List> finalVectors = vectors;
+ final List> finalQueries = queries;
+
+ System.out.format("\t%d base and %d query vectors loaded, dimension=%d%n",
+ nVectors, nQueries, dimension);
+
+ // ------------------------------------------------------------
+ // Parallel executor (identical to DistancesASH)
+ // ------------------------------------------------------------
+ ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool();
+ int parallelism = simdExecutor.getParallelism();
+ int chunkSize = Math.max(1, (nQueries + parallelism - 1) / parallelism);
+
+ // ------------------------------------------------------------
+ // PQ parameters (compression matched to ASH)
+ // ------------------------------------------------------------
+ // ASH: encodedBits = d / 4 -> bytes = d / 32
+ // PQ: 1 byte per subspace -> M = d / 32
+ final int M = Math.max(1, dimension / 32);
+ final int K = 256; // 1 byte per subspace
+ final boolean globallyCenter = false;
+ final float anisotropicThreshold = 0.0f;
+
+ System.out.println("\tPQ params:");
+ System.out.println("\t M = " + M + " subspaces");
+ System.out.println("\t K = " + K + " clusters");
+ System.out.println("\t Compression ≈ " + (M * 8) + " bits (" + M + " bytes)");
+ System.out.println("\t ASH reference ≈ " + (dimension / 4) + " bits (" + (dimension / 32) + " bytes)");
+
+ // ------------------------------------------------------------
+ // Train PQ + encode (timed)
+ // ------------------------------------------------------------
+ RandomAccessVectorValues ravv =
+ new ListRandomAccessVectorValues(finalVectors, dimension);
+
+ long pqTrainStart = System.nanoTime();
+ ProductQuantization pq = ProductQuantization.compute(
+ ravv,
+ M,
+ K,
+ globallyCenter,
+ anisotropicThreshold,
+ simdExecutor,
+ ForkJoinPool.commonPool()
+ );
+ long pqTrainEnd = System.nanoTime();
+
+ long pqEncodeStart = System.nanoTime();
+ CompressedVectors pqAny = pq.encodeAll(ravv, simdExecutor);
+ long pqEncodeEnd = System.nanoTime();
+
+ PQVectors pqVecs = (PQVectors) pqAny;
+
+ System.out.println("\tPQ training took " + (pqTrainEnd - pqTrainStart) / 1e9 + " seconds");
+ System.out.println("\tPQ encoding took " + (pqEncodeEnd - pqEncodeStart) / 1e9 + " seconds");
+
+ // ============================================================
+ // [1] Accuracy (NOT timed)
+ // Compare unscaled dot products:
+ // PQ returns (1 + dp)/2 -> dp = 2*score - 1
+ // ============================================================
+ List> errorTasks = new ArrayList<>();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ errorTasks.add(simdExecutor.submit(() -> {
+ double localError = 0.0;
+ long localCount = 0;
+
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ ScoreFunction.ApproximateScoreFunction f =
+ pqVecs.precomputedScoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+
+ for (int j = 0; j < nVectors; j++) {
+ float trueDot = VectorUtil.dotProduct(q, finalVectors.get(j));
+ float approxScaled = f.similarityTo(j);
+ float approxDot = 2.0f * approxScaled - 1.0f;
+
+ localError += Math.abs(approxDot - trueDot);
+ localCount++;
+ }
+ }
+ return new double[]{localError, localCount};
+ }));
+ }
+
+ double totalError = 0.0;
+ long totalCount = 0;
+ for (ForkJoinTask t : errorTasks) {
+ double[] r = t.join();
+ totalError += r[0];
+ totalCount += (long) r[1];
+ }
+
+ System.out.println("\tAverage absolute dot-product error = " + (totalError / totalCount));
+
+ // ============================================================
+ // [2] Fast PQ scan timing (PQDecoder / ADC)
+ // ============================================================
+ List> pqTasks = new ArrayList<>();
+ long pqStart = System.nanoTime();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ pqTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
+ for (int i = s; i < e; i++) {
+ ScoreFunction.ApproximateScoreFunction f =
+ pqVecs.precomputedScoreFunctionFor(
+ finalQueries.get(i),
+ VectorSimilarityFunction.DOT_PRODUCT
+ );
+ for (int j = 0; j < nVectors; j++) {
+ localSum += f.similarityTo(j);
+ }
+ }
+ return localSum;
+ }));
+ }
+
+ double pqDummy = 0.0;
+ for (ForkJoinTask t : pqTasks) {
+ pqDummy += t.join();
+ }
+
+ long pqEnd = System.nanoTime();
+ System.out.println("\tPQDecoder scan took " + (pqEnd - pqStart) / 1e9 + " seconds");
+
+ // ============================================================
+ // [3] Float dot-product baseline
+ // ============================================================
+ List> floatTasks = new ArrayList<>();
+ long floatStart = System.nanoTime();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ floatTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ for (int j = 0; j < nVectors; j++) {
+ localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
+ }
+ }
+ return localSum;
+ }));
+ }
+
+ double floatDummy = 0.0;
+ for (ForkJoinTask t : floatTasks) {
+ floatDummy += t.join();
+ }
+
+ long floatEnd = System.nanoTime();
+ System.out.println("\tFloat dot-product scan took " + (floatEnd - floatStart) / 1e9 + " seconds");
+
+ // Prevent dead-code elimination
+ System.out.println("\tdummyAccumulator = " + (float) (pqDummy + floatDummy));
+ System.out.println("--");
+ }
+
+ // ------------------------------------------------------------------
+ // Dataset runners (mirrors DistancesASH)
+ // ------------------------------------------------------------------
+
+ public static void runCohere100k() throws IOException {
+ System.out.println("Running Cohere-100k");
+ testPQEncodings(
+ "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_base_vectors_100000.fvec",
+ "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_query_vectors_10000.fvec"
+ );
+ }
+
+ public static void runADA() throws IOException {
+ System.out.println("Running ada_002");
+ testPQEncodings(
+ "./fvec/ada-002/ada_002_100000_base_vectors.fvec",
+ "./fvec/ada-002/ada_002_100000_query_vectors_10000.fvec"
+ );
+ }
+
+ public static void runOpenai1536() throws IOException {
+ System.out.println("Running text-embedding-3-large_1536");
+ testPQEncodings(
+ "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_base_vectors.fvec",
+ "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_query_vectors_10000.fvec"
+ );
+ }
+
+ public static void runOpenai3072() throws IOException {
+ System.out.println("Running text-embedding-3-large_3072");
+ testPQEncodings(
+ "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_base_vectors.fvec",
+ "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_query_vectors_10000.fvec"
+ );
+ }
+
+ public static void main(String[] args) throws IOException {
+ runCohere100k();
+ runADA();
+ runOpenai1536();
+ runOpenai3072();
+ }
+}
From d3dd57b074344311301d3d49ba08f9ffcdc01bde Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Thu, 1 Jan 2026 00:31:48 +0000
Subject: [PATCH 010/243] Added benchmarking options to DistancesASH and
DistancesPQ. Introduced a PQ entropy analysis program.
---
.../jbellis/jvector/example/DistancesASH.java | 264 ++++++++++--------
.../jbellis/jvector/example/DistancesPQ.java | 164 ++++++-----
.../jvector/example/PQEntropyAnalysis.java | 198 +++++++++++++
3 files changed, 449 insertions(+), 177 deletions(-)
create mode 100644 jvector-examples/src/main/java/io/github/jbellis/jvector/example/PQEntropyAnalysis.java
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 0fac759fe..2d6797bab 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -48,8 +48,23 @@ private static void printScorerInfo(String label, ASHBlockScorer scorer) {
}
public static void testASHEncodings(String filenameBase, String filenameQueries) throws IOException {
+ // ------------------------------------------------------------
+ // Benchmark configuration (runtime flags)
+ // ------------------------------------------------------------
+ final boolean RUN_SANITY_CHECK =
+ Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
+
+ final boolean RUN_ACCURACY_CHECK =
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
+
+ final boolean RUN_SCALAR_SCORING =
+ Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
+
+ final boolean RUN_FLOAT_SCORING =
+ Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
+
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {16, 32, 64, 128}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {64}; // 16, 32, and/or 64
List> vectors = SiftLoader.readFvecs(filenameBase);
List> queries = SiftLoader.readFvecs(filenameQueries);
@@ -97,8 +112,8 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
", encodedBits = " + encodedBits
);
- int nQueries = Math.min(1000, queries.size());
- int nVectors = Math.min(100_000, vectors.size());
+ int nQueries = Math.min(10000, queries.size());
+ int nVectors = Math.min(1_000_000, vectors.size());
vectors = vectors.subList(0, nVectors);
queries = queries.subList(0, nQueries);
@@ -123,37 +138,40 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
ASHVectors ashVectors = (ASHVectors) ashVecs;
VectorFloat> q0 = finalQueries.get(0);
- ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
- printScorerInfo("scalar-ref", ref);
+ if(RUN_SANITY_CHECK) {
- final int sanityBlockSize = BLOCK_SIZES[0];
- ASHBlockScorer blk = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT, sanityBlockSize);
- printScorerInfo("block", blk);
+ ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
+ printScorerInfo("scalar-ref", ref);
+ final int sanityBlockSize = BLOCK_SIZES[0];
+ ASHBlockScorer blk = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT, sanityBlockSize);
+ printScorerInfo("block", blk);
- float[] a = new float[128];
- float[] b = new float[128];
- ref.scoreRange(0, 128, a);
- blk.scoreRange(0, 128, b);
+ float[] a = new float[128];
+ float[] b = new float[128];
+ ref.scoreRange(0, 128, a);
+ blk.scoreRange(0, 128, b);
- final float EPS = 1e-6f;
+ final float EPS = 1e-6f;
- for (int i = 0; i < 128; i++) {
- float diff = Math.abs(a[i] - b[i]);
- if (diff > EPS) {
- throw new AssertionError(
- "Mismatch at " + i +
- ": ref=" + a[i] +
- ", blk=" + b[i] +
- ", diff=" + diff);
+ for (int i = 0; i < 128; i++) {
+ float diff = Math.abs(a[i] - b[i]);
+ if (diff > EPS) {
+ throw new AssertionError(
+ "Mismatch at " + i +
+ ": ref=" + a[i] +
+ ", blk=" + b[i] +
+ ", diff=" + diff);
+ }
}
}
// ------------------------------------------------------------------
// Shared parallelization setup (executor-honoring)
// ------------------------------------------------------------------
- ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool();
+ // ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool(); // For production
+ ForkJoinPool simdExecutor = new ForkJoinPool(4); // For profiling
int parallelism = simdExecutor.getParallelism();
int chunkSize = Math.max(1, (nQueries + parallelism - 1) / parallelism);
@@ -161,120 +179,128 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// ==================================================================
// [1] Accuracy run (NOT timed)
// ==================================================================
- List> errorTasks = new ArrayList<>();
+ if(RUN_ACCURACY_CHECK) {
+ List> errorTasks = new ArrayList<>();
- for (int start = 0; start < nQueries; start += chunkSize) {
- final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
- errorTasks.add(simdExecutor.submit(() -> {
- double localError = 0.0;
- long localCount = 0;
+ errorTasks.add(simdExecutor.submit(() -> {
+ double localError = 0.0;
+ long localCount = 0;
- for (int i = s; i < e; i++) {
- VectorFloat> q = finalQueries.get(i);
- ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ ScoreFunction.ApproximateScoreFunction f =
+ ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
- for (int j = 0; j < nVectors; j++) {
- VectorFloat> v = finalVectors.get(j);
- float trueDot = VectorUtil.dotProduct(q, v);
- float approxDot = f.similarityTo(j);
+ for (int j = 0; j < nVectors; j++) {
+ VectorFloat> v = finalVectors.get(j);
+ float trueDot = VectorUtil.dotProduct(q, v);
+ float approxDot = f.similarityTo(j);
- localError += Math.abs(approxDot - trueDot);
- localCount++;
+ localError += Math.abs(approxDot - trueDot);
+ localCount++;
+ }
}
- }
- return new double[]{localError, localCount};
- }));
- }
+ return new double[]{localError, localCount};
+ }));
+ }
- double distanceError = 0.0;
- long count = 0;
+ double distanceError = 0.0;
+ long count = 0;
- for (ForkJoinTask t : errorTasks) {
- double[] r = t.join();
- distanceError += r[0];
- count += (long) r[1];
- }
+ for (ForkJoinTask t : errorTasks) {
+ double[] r = t.join();
+ distanceError += r[0];
+ count += (long) r[1];
+ }
- distanceError /= count;
- System.out.println("\tAverage absolute dot-product error = " + distanceError);
+ distanceError /= count;
+ System.out.println("\tAverage absolute dot-product error = " + distanceError);
+ }
// ==================================================================
- // [2] Scalar ASH distance timing (reference)
+ // [2] Scalar ASH scoring timing (reference)
// ==================================================================
- List> ashTasks = new ArrayList<>();
+ if (RUN_SCALAR_SCORING) {
+ List> ashTasks = new ArrayList<>();
- long ashStart = System.nanoTime();
+ long ashStart = System.nanoTime();
- for (int start = 0; start < nQueries; start += chunkSize) {
- final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
- ashTasks.add(simdExecutor.submit(() -> {
- double localSum = 0.0;
+ ashTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
- for (int i = s; i < e; i++) {
- VectorFloat> q = finalQueries.get(i);
- ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ ScoreFunction.ApproximateScoreFunction f =
+ ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
- for (int j = 0; j < nVectors; j++) {
- localSum += f.similarityTo(j);
+ for (int j = 0; j < nVectors; j++) {
+ localSum += f.similarityTo(j);
+ }
}
- }
- return localSum;
- }));
- }
+ return localSum;
+ }));
+ }
- double ashDummy = 0.0;
- for (ForkJoinTask t : ashTasks) {
- ashDummy += t.join();
- }
+ double ashDummy = 0.0;
+ for (ForkJoinTask t : ashTasks) {
+ ashDummy += t.join();
+ }
- long ashEnd = System.nanoTime();
+ long ashEnd = System.nanoTime();
- System.out.println("\tScalar ASH dot-product computations took "
- + (ashEnd - ashStart) / 1e9 + " seconds");
+ System.out.println("\tScalar ASH dot-product computations took "
+ + (ashEnd - ashStart) / 1e9 + " seconds");
+ System.out.println("\tdummyAccumulator = " + (float) (ashDummy));
+ System.out.println("--");
+ }
// ==================================================================
- // [3] Float dot-product timing (ground truth baseline)
+ // [3] Float dot-product scoring (ground truth baseline)
// ==================================================================
- List> floatTasks = new ArrayList<>();
+ if (RUN_FLOAT_SCORING) {
+ List> floatTasks = new ArrayList<>();
- long floatStart = System.nanoTime();
+ long floatStart = System.nanoTime();
- for (int start = 0; start < nQueries; start += chunkSize) {
- final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
- floatTasks.add(simdExecutor.submit(() -> {
- double localSum = 0.0;
+ floatTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
- for (int i = s; i < e; i++) {
- VectorFloat> q = finalQueries.get(i);
- for (int j = 0; j < nVectors; j++) {
- localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ for (int j = 0; j < nVectors; j++) {
+ localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
+ }
}
- }
- return localSum;
- }));
- }
+ return localSum;
+ }));
+ }
- double floatDummy = 0.0;
- for (ForkJoinTask t : floatTasks) {
- floatDummy += t.join();
- }
+ double floatDummy = 0.0;
+ for (ForkJoinTask t : floatTasks) {
+ floatDummy += t.join();
+ }
- long floatEnd = System.nanoTime();
+ long floatEnd = System.nanoTime();
- System.out.println("\tFloat dot-product computations took "
- + (floatEnd - floatStart) / 1e9 + " seconds");
+ System.out.println("\tFloat dot-product computations took "
+ + (floatEnd - floatStart) / 1e9 + " seconds");
- // Prevent dead-code elimination
- System.out.println("\tdummyAccumulator = " + (float) (ashDummy + floatDummy));
- System.out.println("--");
+ // Prevent dead-code elimination
+ System.out.println("\tdummyAccumulator = " + (float) (floatDummy));
+ System.out.println("--");
+ }
// ==================================================================
// [4] Block scoring timing (single mode per JVM run)
@@ -355,11 +381,9 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
+ (blockEnd - blockStart) / 1e9 + " seconds"
);
- ashDummy += blockDummy;
+ System.out.println("\tdummyAccumulator = " + (float) (blockDummy));
+ System.out.println("--");
}
- // Prevent dead-code elimination across all sections
- System.out.println("\tdummyAccumulator = " + (float) (ashDummy + floatDummy));
- System.out.println("--");
}
public static void runSIFT() throws IOException {
@@ -418,13 +442,31 @@ public static void runOpenai3072() throws IOException {
testASHEncodings(baseVectors, queryVectors);
}
+ public static void runCap6m() throws IOException {
+ System.out.println("Running cap-6m");
+
+ var baseVectors = "./fvec/cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_base_6m_norm_shuffle.fvecs";
+ var queryVectors = "./fvec/cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_query_10k_norm_shuffle.fvecs";
+ testASHEncodings(baseVectors, queryVectors);
+ }
+
+ public static void runCohere10m() throws IOException {
+ System.out.println("Running cohere-10m");
+
+ var baseVectors = "./fvec/cohere-10m/cohere_wiki_en_flat_base_10m_norm.fvecs";
+ var queryVectors = "./fvec/cohere-10m/cohere_wiki_en_flat_query_10k_norm.fvecs";
+ testASHEncodings(baseVectors, queryVectors);
+ }
+
public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
- runCohere100k();
- runADA();
- runOpenai1536();
- runOpenai3072();
+// runCohere100k();
+// runADA();
+// runOpenai1536();
+// runOpenai3072();
+// runCap6m();
+ runCohere10m();
}
}
\ No newline at end of file
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
index e57727fa9..570e663f4 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -48,6 +48,13 @@ public class DistancesPQ {
* Fast PQ path = PQDecoder (precomputedScoreFunctionFor).
*/
public static void testPQEncodings(String filenameBase, String filenameQueries) throws IOException {
+ // ------------------------------------------------------------
+ // Benchmark configuration (runtime flags)
+ // ------------------------------------------------------------
+ final boolean RUN_ACCURACY_CHECK =
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy-check", "true"));
+ final boolean RUN_FLOAT_SCORING =
+ Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
List> vectors = SiftLoader.readFvecs(filenameBase);
List> queries = SiftLoader.readFvecs(filenameQueries);
@@ -83,8 +90,8 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
int dimension = vectors.get(0).length();
- int nQueries = Math.min(1000, queries.size());
- int nVectors = Math.min(100_000, vectors.size());
+ int nQueries = Math.min(10000, queries.size());
+ int nVectors = Math.min(10_000_000, vectors.size());
vectors = vectors.subList(0, nVectors);
queries = queries.subList(0, nQueries);
@@ -150,43 +157,45 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
// Compare unscaled dot products:
// PQ returns (1 + dp)/2 -> dp = 2*score - 1
// ============================================================
- List> errorTasks = new ArrayList<>();
-
- for (int start = 0; start < nQueries; start += chunkSize) {
- final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
-
- errorTasks.add(simdExecutor.submit(() -> {
- double localError = 0.0;
- long localCount = 0;
-
- for (int i = s; i < e; i++) {
- VectorFloat> q = finalQueries.get(i);
- ScoreFunction.ApproximateScoreFunction f =
- pqVecs.precomputedScoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
-
- for (int j = 0; j < nVectors; j++) {
- float trueDot = VectorUtil.dotProduct(q, finalVectors.get(j));
- float approxScaled = f.similarityTo(j);
- float approxDot = 2.0f * approxScaled - 1.0f;
-
- localError += Math.abs(approxDot - trueDot);
- localCount++;
+ if (RUN_ACCURACY_CHECK) {
+ List> errorTasks = new ArrayList<>();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ errorTasks.add(simdExecutor.submit(() -> {
+ double localError = 0.0;
+ long localCount = 0;
+
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ ScoreFunction.ApproximateScoreFunction f =
+ pqVecs.precomputedScoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+
+ for (int j = 0; j < nVectors; j++) {
+ float trueDot = VectorUtil.dotProduct(q, finalVectors.get(j));
+ float approxScaled = f.similarityTo(j);
+ float approxDot = 2.0f * approxScaled - 1.0f;
+
+ localError += Math.abs(approxDot - trueDot);
+ localCount++;
+ }
}
- }
- return new double[]{localError, localCount};
- }));
- }
+ return new double[]{localError, localCount};
+ }));
+ }
- double totalError = 0.0;
- long totalCount = 0;
- for (ForkJoinTask t : errorTasks) {
- double[] r = t.join();
- totalError += r[0];
- totalCount += (long) r[1];
- }
+ double totalError = 0.0;
+ long totalCount = 0;
+ for (ForkJoinTask t : errorTasks) {
+ double[] r = t.join();
+ totalError += r[0];
+ totalCount += (long) r[1];
+ }
- System.out.println("\tAverage absolute dot-product error = " + (totalError / totalCount));
+ System.out.println("\tAverage absolute dot-product error = " + (totalError / totalCount));
+ }
// ============================================================
// [2] Fast PQ scan timing (PQDecoder / ADC)
@@ -218,6 +227,9 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
for (ForkJoinTask t : pqTasks) {
pqDummy += t.join();
}
+ // Prevent dead-code elimination
+ System.out.println("\tdummyAccumulator = " + (float) (pqDummy));
+ System.out.println("--");
long pqEnd = System.nanoTime();
System.out.println("\tPQDecoder scan took " + (pqEnd - pqStart) / 1e9 + " seconds");
@@ -225,36 +237,38 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
// ============================================================
// [3] Float dot-product baseline
// ============================================================
- List> floatTasks = new ArrayList<>();
- long floatStart = System.nanoTime();
-
- for (int start = 0; start < nQueries; start += chunkSize) {
- final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
-
- floatTasks.add(simdExecutor.submit(() -> {
- double localSum = 0.0;
- for (int i = s; i < e; i++) {
- VectorFloat> q = finalQueries.get(i);
- for (int j = 0; j < nVectors; j++) {
- localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
+ if (RUN_FLOAT_SCORING) {
+ List> floatTasks = new ArrayList<>();
+ long floatStart = System.nanoTime();
+
+ for (int start = 0; start < nQueries; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, nQueries);
+
+ floatTasks.add(simdExecutor.submit(() -> {
+ double localSum = 0.0;
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ for (int j = 0; j < nVectors; j++) {
+ localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
+ }
}
- }
- return localSum;
- }));
- }
+ return localSum;
+ }));
+ }
- double floatDummy = 0.0;
- for (ForkJoinTask t : floatTasks) {
- floatDummy += t.join();
- }
+ double floatDummy = 0.0;
+ for (ForkJoinTask t : floatTasks) {
+ floatDummy += t.join();
+ }
- long floatEnd = System.nanoTime();
- System.out.println("\tFloat dot-product scan took " + (floatEnd - floatStart) / 1e9 + " seconds");
+ long floatEnd = System.nanoTime();
+ System.out.println("\tFloat dot-product scan took " + (floatEnd - floatStart) / 1e9 + " seconds");
- // Prevent dead-code elimination
- System.out.println("\tdummyAccumulator = " + (float) (pqDummy + floatDummy));
- System.out.println("--");
+ // Prevent dead-code elimination
+ System.out.println("\tdummyAccumulator = " + (float) (floatDummy));
+ System.out.println("--");
+ }
}
// ------------------------------------------------------------------
@@ -293,10 +307,28 @@ public static void runOpenai3072() throws IOException {
);
}
+ public static void runCap6m() throws IOException {
+ System.out.println("Running cap-6m");
+
+ var baseVectors = "./fvec/cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_base_6m_norm_shuffle.fvecs";
+ var queryVectors = "./fvec/cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_query_10k_norm_shuffle.fvecs";
+ testPQEncodings(baseVectors, queryVectors);
+ }
+
+ public static void runCohere10m() throws IOException {
+ System.out.println("Running cohere-10m");
+
+ var baseVectors = "./fvec/cohere-10m/cohere_wiki_en_flat_base_10m_norm.fvecs";
+ var queryVectors = "./fvec/cohere-10m/cohere_wiki_en_flat_query_10k_norm.fvecs";
+ testPQEncodings(baseVectors, queryVectors);
+ }
+
public static void main(String[] args) throws IOException {
- runCohere100k();
- runADA();
- runOpenai1536();
- runOpenai3072();
+// runCohere100k();
+// runADA();
+// runOpenai1536();
+// runOpenai3072();
+ runCap6m();
+ runCohere10m();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/PQEntropyAnalysis.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/PQEntropyAnalysis.java
new file mode 100644
index 000000000..9b5aa9d77
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/PQEntropyAnalysis.java
@@ -0,0 +1,198 @@
+/*
+ * 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;
+
+import io.github.jbellis.jvector.example.util.SiftLoader;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
+import io.github.jbellis.jvector.quantization.PQVectors;
+import io.github.jbellis.jvector.quantization.ProductQuantization;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.ForkJoinPool;
+
+/**
+ * Diagnostic tool to analyze Product Quantization centroid assignment entropy.
+ *
+ *
+ * Explains PQ scan-time behavior by measuring how uniformly vectors are
+ * assigned to centroids in each subspace.
+ *
+ *
+ *
+ * High entropy ⇒ many centroids exercised ⇒ poor cache locality and slower ADC.
+ * Low entropy ⇒ concentrated centroids ⇒ faster PQ scans.
+ *
+ *
+ *
+ * This is NOT a benchmark:
+ *
+ * - No timing loops
+ * - No scoring
+ * - No SIMD assumptions
+ *
+ *
+ */
+public class PQEntropyAnalysis {
+
+ public static void analyze(
+ String label,
+ String baseVectorsPath,
+ int encodedBits
+ ) throws IOException {
+
+ System.out.println("--------------------------------------------------");
+ System.out.println("PQ Entropy Analysis: " + label);
+
+ // ------------------------------------------------------------
+ // Load base vectors
+ // ------------------------------------------------------------
+ List> vectors = SiftLoader.readFvecs(baseVectorsPath);
+ int dimension = vectors.get(0).length();
+
+ System.out.println("Vectors: " + vectors.size());
+ System.out.println("Dimension: " + dimension);
+ System.out.println("Target compression: " + encodedBits + " bits");
+
+ // ------------------------------------------------------------
+ // PQ parameters (match DistancesPQ)
+ // ------------------------------------------------------------
+ // ASH: encodedBits = d / 4 → bytes = d / 32
+ // PQ: 1 byte per subspace → M = d / 32
+ final int M = Math.max(1, encodedBits / 8);
+ final int K = 256;
+ final boolean globallyCenter = false;
+ final float anisotropicThreshold = 0.0f;
+
+ System.out.println("PQ params:");
+ System.out.println(" M (subspaces) = " + M);
+ System.out.println(" K (clusters) = " + K);
+
+ // ------------------------------------------------------------
+ // Build PQ
+ // ------------------------------------------------------------
+ RandomAccessVectorValues ravv =
+ new ListRandomAccessVectorValues(vectors, dimension);
+
+ ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool();
+
+ ProductQuantization pq = ProductQuantization.compute(
+ ravv,
+ M,
+ K,
+ globallyCenter,
+ anisotropicThreshold,
+ simdExecutor,
+ ForkJoinPool.commonPool()
+ );
+
+ PQVectors pqVecs = pq.encodeAll(ravv, simdExecutor);
+
+ // ------------------------------------------------------------
+ // Count centroid assignments
+ // ------------------------------------------------------------
+ long[][] counts = new long[M][K];
+
+ for (int i = 0; i < pqVecs.count(); i++) {
+ var code = pqVecs.get(i);
+ for (int m = 0; m < M; m++) {
+ int c = Byte.toUnsignedInt(code.get(m));
+ counts[m][c]++;
+ }
+ }
+
+ // ------------------------------------------------------------
+ // Compute entropy per subspace
+ // ------------------------------------------------------------
+ double totalEntropy = 0.0;
+
+ for (int m = 0; m < M; m++) {
+ long total = 0;
+ for (int k = 0; k < K; k++) {
+ total += counts[m][k];
+ }
+
+ double entropy = 0.0;
+ for (int k = 0; k < K; k++) {
+ long cnt = counts[m][k];
+ if (cnt == 0) continue;
+ double p = (double) cnt / total;
+ entropy -= p * (Math.log(p) / Math.log(2));
+ }
+
+ totalEntropy += entropy;
+ }
+
+ double avgEntropy = totalEntropy / M;
+ double effectiveCentroids = Math.pow(2.0, avgEntropy);
+
+ // ------------------------------------------------------------
+ // Report
+ // ------------------------------------------------------------
+ System.out.printf("Subspaces (M): %d%n", M);
+ System.out.printf("Centroids per subspace (K): %d%n", K);
+ System.out.printf("Average entropy per subspace: %.4f bits%n", avgEntropy);
+ System.out.printf("Effective centroids used: %.1f%n", effectiveCentroids);
+ }
+
+ // ------------------------------------------------------------------
+ // Dataset runners (aligned with DistancesPQ / DistancesASH)
+ // ------------------------------------------------------------------
+
+ public static void runCohere100k() throws IOException {
+ analyze(
+ "Cohere-100k",
+ "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_base_vectors_100000.fvec",
+ 256
+ );
+ }
+
+ public static void runAda002() throws IOException {
+ analyze(
+ "ada_002",
+ "./fvec/ada-002/ada_002_100000_base_vectors.fvec",
+ 384
+ );
+ }
+
+ public static void runOpenAI1536() throws IOException {
+ analyze(
+ "text-embedding-3-large-1536",
+ "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_base_vectors.fvec",
+ 384
+ );
+ }
+
+ public static void runOpenAI3072() throws IOException {
+ analyze(
+ "text-embedding-3-large-3072",
+ "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_base_vectors.fvec",
+ 768
+ );
+ }
+
+ public static void main(String[] args) throws IOException {
+ runCohere100k();
+ runAda002();
+ runOpenAI1536();
+ runOpenAI3072();
+ }
+}
+
From 22c4f53eb6b2ad11cedfbe9ce2831cce59454505 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Thu, 1 Jan 2026 02:32:00 +0000
Subject: [PATCH 011/243] We now amortize the cost of horizontal reductions
during ASH scoring.
---
.../jbellis/jvector/example/DistancesASH.java | 20 +++++++++----------
.../vector/PanamaVectorUtilSupport.java | 20 +++++++++++++++----
2 files changed, 26 insertions(+), 14 deletions(-)
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 2d6797bab..bf5438810 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -55,7 +55,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "true"));
final boolean RUN_SCALAR_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
@@ -64,7 +64,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {64}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {16, 32, 64, 128}; // 16, 32, and/or 64
List> vectors = SiftLoader.readFvecs(filenameBase);
List> queries = SiftLoader.readFvecs(filenameQueries);
@@ -112,7 +112,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
", encodedBits = " + encodedBits
);
- int nQueries = Math.min(10000, queries.size());
+ int nQueries = Math.min(1000, queries.size());
int nVectors = Math.min(1_000_000, vectors.size());
vectors = vectors.subList(0, nVectors);
@@ -170,8 +170,8 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// ------------------------------------------------------------------
// Shared parallelization setup (executor-honoring)
// ------------------------------------------------------------------
- // ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool(); // For production
- ForkJoinPool simdExecutor = new ForkJoinPool(4); // For profiling
+ ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool(); // For production
+ // ForkJoinPool simdExecutor = new ForkJoinPool(4); // For profiling
int parallelism = simdExecutor.getParallelism();
int chunkSize = Math.max(1, (nQueries + parallelism - 1) / parallelism);
@@ -462,11 +462,11 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
-// runCohere100k();
-// runADA();
-// runOpenai1536();
-// runOpenai3072();
+ runCohere100k();
+ runADA();
+ runOpenai1536();
+ runOpenai3072();
// runCap6m();
- runCohere10m();
+// runCohere10m();
}
}
\ No newline at end of file
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 fbd9483eb..33bab6aea 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
@@ -1591,11 +1591,20 @@ public void ashMaskedAddBlockAllWords(float[] tildeQ,
// Like in the ASH paper, SIMD lanes are dimensions (SPECIES_PREFERRED),
// masked load from tildeQ and horizontal sum per neighbor.
+ //
+ // Optimization: amortize horizontal reductions.
+ // Instead of reducing each masked chunk immediately, we accumulate masked
+ // vectors into a vector register accumulator and reduce once per lane.
final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
final int LANES = SPEC.length();
for (int lane = 0; lane < blockLen; lane++) {
- float maskedAdd = 0f; // register scalar accumulator
+ // Vector accumulator: stays in registers (no array accumulation).
+ FloatVector accVec = FloatVector.zero(SPEC);
+
+ // Scalar accumulator for unusual widths fallback only.
+ float scalarAcc = 0f;
+
int laneIndex = laneStart + lane;
for (int w = 0; w < words; w++) {
@@ -1626,7 +1635,7 @@ public void ashMaskedAddBlockAllWords(float[] tildeQ,
if (idx < d) s += tildeQ[idx];
}
}
- maskedAdd += s;
+ scalarAcc += s;
continue;
}
@@ -1638,11 +1647,14 @@ public void ashMaskedAddBlockAllWords(float[] tildeQ,
// masked load (paper)
FloatVector masked = FloatVector.fromArray(SPEC, tildeQ, qBase, bitMask);
- // horizontal sum (paper)
- maskedAdd += masked.reduceLanes(VectorOperators.ADD);
+ // amortized accumulation (register)
+ accVec = accVec.add(masked);
}
}
+ // One horizontal sum per lane (instead of per chunk)
+ float maskedAdd = accVec.reduceLanes(VectorOperators.ADD) + scalarAcc;
+
outMaskedAdd[lane] = maskedAdd;
}
}
From 2d8c6fd20e161e7d01a38741130c6ab6e4bf6456 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Thu, 1 Jan 2026 15:21:48 -0800
Subject: [PATCH 012/243] Added overload to SiftLoader to allow for truncation
of the load (without loading all vectors first). Updated DistancesASH and
DistancesPQ to take advantage of this. The original SiftLoader remains
intact.
---
.../jbellis/jvector/example/DistancesASH.java | 74 ++++----
.../jbellis/jvector/example/DistancesPQ.java | 61 ++++---
.../jvector/example/util/SiftLoader.java | 161 ++++++++++++++++--
3 files changed, 223 insertions(+), 73 deletions(-)
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index bf5438810..ac8fe886a 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -32,13 +32,10 @@
import java.io.IOException;
import java.util.List;
-import java.util.stream.Collectors;
import java.util.ArrayList;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;
-import static java.lang.Math.abs;
-
public class DistancesASH {
private static void printScorerInfo(String label, ASHBlockScorer scorer) {
System.out.println(
@@ -64,10 +61,21 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {16, 32, 64, 128}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {32}; // 16, 32, and/or 64
+
+ // Define the benchmark size
+ int maxQueries = 10_000;
+ int maxVectors = 3_000_000;
+
+ int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
+ int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
- List> vectors = SiftLoader.readFvecs(filenameBase);
- List> queries = SiftLoader.readFvecs(filenameQueries);
+ // Will benchmark what was requested or the maximum available vectors (if less)
+ int nQueries = Math.min(maxQueries, queryCountInFile);
+ int nVectors = Math.min(maxVectors, vectorCountInFile);
+
+ List> vectors = SiftLoader.readFvecs(filenameBase, nVectors);
+ List> queries = SiftLoader.readFvecs(filenameQueries, nQueries);
// ------------------------------------------------------------------
// Remove zero-norm vectors (undefined for ASH / Eq. 11)
@@ -93,6 +101,12 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
System.out.println("\tRemoved " + (beforeBase - vectors.size()) + " zero base vectors and "
+ (beforeQuery - queries.size()) + " zero query vectors");
+ // ------------------------------------------------------------
+ // Update counts after filtering
+ // ------------------------------------------------------------
+ final int finalVectorCount = vectors.size();
+ final int finalQueryCount = queries.size();
+
// ASH normalization policy:
//
// - Base vectors x are NOT globally normalized.
@@ -112,12 +126,6 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
", encodedBits = " + encodedBits
);
- int nQueries = Math.min(1000, queries.size());
- int nVectors = Math.min(1_000_000, vectors.size());
-
- vectors = vectors.subList(0, nVectors);
- queries = queries.subList(0, nQueries);
-
final List> finalQueries = queries;
final List> finalVectors = vectors;
@@ -174,7 +182,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// ForkJoinPool simdExecutor = new ForkJoinPool(4); // For profiling
int parallelism = simdExecutor.getParallelism();
- int chunkSize = Math.max(1, (nQueries + parallelism - 1) / parallelism);
+ int chunkSize = Math.max(1, (finalQueryCount + parallelism - 1) / parallelism);
// ==================================================================
// [1] Accuracy run (NOT timed)
@@ -182,9 +190,9 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
if(RUN_ACCURACY_CHECK) {
List> errorTasks = new ArrayList<>();
- for (int start = 0; start < nQueries; start += chunkSize) {
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ final int e = Math.min(start + chunkSize, finalQueryCount);
errorTasks.add(simdExecutor.submit(() -> {
double localError = 0.0;
@@ -195,7 +203,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
ScoreFunction.ApproximateScoreFunction f =
ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
- for (int j = 0; j < nVectors; j++) {
+ for (int j = 0; j < finalVectorCount; j++) {
VectorFloat> v = finalVectors.get(j);
float trueDot = VectorUtil.dotProduct(q, v);
float approxDot = f.similarityTo(j);
@@ -229,9 +237,9 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
long ashStart = System.nanoTime();
- for (int start = 0; start < nQueries; start += chunkSize) {
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ final int e = Math.min(start + chunkSize, finalQueryCount);
ashTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
@@ -241,7 +249,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
ScoreFunction.ApproximateScoreFunction f =
ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
- for (int j = 0; j < nVectors; j++) {
+ for (int j = 0; j < finalVectorCount; j++) {
localSum += f.similarityTo(j);
}
}
@@ -270,16 +278,16 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
long floatStart = System.nanoTime();
- for (int start = 0; start < nQueries; start += chunkSize) {
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ final int e = Math.min(start + chunkSize, finalQueryCount);
floatTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
for (int i = s; i < e; i++) {
VectorFloat> q = finalQueries.get(i);
- for (int j = 0; j < nVectors; j++) {
+ for (int j = 0; j < finalVectorCount; j++) {
localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
}
}
@@ -327,9 +335,9 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
List> blockTasks = new ArrayList<>();
long blockStart = System.nanoTime();
- for (int start = 0; start < nQueries; start += chunkSize) {
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ final int e = Math.min(start + chunkSize, finalQueryCount);
blockTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
@@ -344,7 +352,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
);
int j = 0;
- while (j + blockSize <= nVectors) {
+ while (j + blockSize <= finalVectorCount) {
scorer.scoreRange(j, blockSize, scores);
for (int k = 0; k < blockSize; k++) {
localSum += scores[k];
@@ -353,13 +361,13 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
}
// Tail uses scalar per-vector scorer (correctness-preserving)
- if (j < nVectors) {
+ if (j < finalVectorCount) {
ScoreFunction.ApproximateScoreFunction f =
ashVecs.scoreFunctionFor(
finalQueries.get(qi),
VectorSimilarityFunction.DOT_PRODUCT
);
- for (; j < nVectors; j++) {
+ for (; j < finalVectorCount; j++) {
localSum += f.similarityTo(j);
}
}
@@ -462,11 +470,11 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
- runCohere100k();
- runADA();
- runOpenai1536();
- runOpenai3072();
+// runCohere100k();
+// runADA();
+// runOpenai1536();
+// runOpenai3072();
// runCap6m();
-// runCohere10m();
+ runCohere10m();
}
-}
\ No newline at end of file
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
index 570e663f4..a11f3c156 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -56,8 +56,19 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
final boolean RUN_FLOAT_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
- List> vectors = SiftLoader.readFvecs(filenameBase);
- List> queries = SiftLoader.readFvecs(filenameQueries);
+ // Define the benchmark size
+ int maxQueries = 10_000;
+ int maxVectors = 2_000_000;
+
+ int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
+ int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
+
+ // Will benchmark what was requested or the maximum available vectors (if less)
+ int nQueries = Math.min(maxQueries, queryCountInFile);
+ int nVectors = Math.min(maxVectors, vectorCountInFile);
+
+ List> vectors = SiftLoader.readFvecs(filenameBase, nVectors);
+ List> queries = SiftLoader.readFvecs(filenameQueries, nQueries);
// ------------------------------------------------------------
// Remove zero-norm vectors (match DistancesASH)
@@ -83,6 +94,12 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
System.out.println("\tRemoved " + (beforeBase - vectors.size()) + " zero base vectors and "
+ (beforeQuery - queries.size()) + " zero query vectors");
+ // ------------------------------------------------------------
+ // Update counts after filtering
+ // ------------------------------------------------------------
+ final int finalVectorCount = vectors.size();
+ final int finalQueryCount = queries.size();
+
// Match ASH: normalize queries only
for (VectorFloat> q : queries) {
VectorUtil.l2normalize(q);
@@ -90,24 +107,18 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
int dimension = vectors.get(0).length();
- int nQueries = Math.min(10000, queries.size());
- int nVectors = Math.min(10_000_000, vectors.size());
-
- vectors = vectors.subList(0, nVectors);
- queries = queries.subList(0, nQueries);
-
final List> finalVectors = vectors;
final List> finalQueries = queries;
System.out.format("\t%d base and %d query vectors loaded, dimension=%d%n",
- nVectors, nQueries, dimension);
+ vectors.size(), queries.size(), dimension);
// ------------------------------------------------------------
// Parallel executor (identical to DistancesASH)
// ------------------------------------------------------------
ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool();
int parallelism = simdExecutor.getParallelism();
- int chunkSize = Math.max(1, (nQueries + parallelism - 1) / parallelism);
+ int chunkSize = Math.max(1, (finalQueryCount + parallelism - 1) / parallelism);
// ------------------------------------------------------------
// PQ parameters (compression matched to ASH)
@@ -160,9 +171,9 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
if (RUN_ACCURACY_CHECK) {
List> errorTasks = new ArrayList<>();
- for (int start = 0; start < nQueries; start += chunkSize) {
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ final int e = Math.min(start + chunkSize, finalQueryCount);
errorTasks.add(simdExecutor.submit(() -> {
double localError = 0.0;
@@ -173,7 +184,7 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
ScoreFunction.ApproximateScoreFunction f =
pqVecs.precomputedScoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
- for (int j = 0; j < nVectors; j++) {
+ for (int j = 0; j < finalVectorCount; j++) {
float trueDot = VectorUtil.dotProduct(q, finalVectors.get(j));
float approxScaled = f.similarityTo(j);
float approxDot = 2.0f * approxScaled - 1.0f;
@@ -203,9 +214,9 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
List> pqTasks = new ArrayList<>();
long pqStart = System.nanoTime();
- for (int start = 0; start < nQueries; start += chunkSize) {
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ final int e = Math.min(start + chunkSize, finalQueryCount);
pqTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
@@ -215,7 +226,7 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
finalQueries.get(i),
VectorSimilarityFunction.DOT_PRODUCT
);
- for (int j = 0; j < nVectors; j++) {
+ for (int j = 0; j < finalVectorCount; j++) {
localSum += f.similarityTo(j);
}
}
@@ -241,15 +252,15 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
List> floatTasks = new ArrayList<>();
long floatStart = System.nanoTime();
- for (int start = 0; start < nQueries; start += chunkSize) {
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
final int s = start;
- final int e = Math.min(start + chunkSize, nQueries);
+ final int e = Math.min(start + chunkSize, finalQueryCount);
floatTasks.add(simdExecutor.submit(() -> {
double localSum = 0.0;
for (int i = s; i < e; i++) {
VectorFloat> q = finalQueries.get(i);
- for (int j = 0; j < nVectors; j++) {
+ for (int j = 0; j < finalVectorCount; j++) {
localSum += VectorUtil.dotProduct(q, finalVectors.get(j));
}
}
@@ -324,11 +335,11 @@ public static void runCohere10m() throws IOException {
}
public static void main(String[] args) throws IOException {
-// runCohere100k();
-// runADA();
-// runOpenai1536();
-// runOpenai3072();
- runCap6m();
- runCohere10m();
+ runCohere100k();
+ runADA();
+ runOpenai1536();
+ runOpenai3072();
+// runCap6m();
+// runCohere10m();
}
}
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 46f144418..cccfc1cea 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
@@ -27,43 +27,137 @@
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.channels.FileChannel;
+import java.nio.file.StandardOpenOption;
import java.util.ArrayList;
-import java.util.HashSet;
import java.util.List;
public class SiftLoader {
- private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+ private static final VectorTypeSupport vectorTypeSupport =
+ VectorizationProvider.getInstance().getVectorTypeSupport();
+
+ /**
+ * Read all vectors from an .fvecs file.
+ *
+ *
+ * Reads the entire file and materializes all vectors into memory.
+ *
+ */
public static List> readFvecs(String filePath) throws IOException {
var vectors = new ArrayList>();
- try (var dis = new DataInputStream(new BufferedInputStream(new FileInputStream(filePath)))) {
+ try (var dis = new DataInputStream(
+ new BufferedInputStream(new FileInputStream(filePath)))) {
+
while (dis.available() > 0) {
- var dimension = Integer.reverseBytes(dis.readInt());
- assert dimension > 0 : dimension;
- var buffer = new byte[dimension * Float.BYTES];
+ int dimension = Integer.reverseBytes(dis.readInt());
+ if (dimension <= 0) {
+ throw new IOException("Invalid vector dimension: " + dimension);
+ }
+
+ byte[] buffer = new byte[dimension * Float.BYTES];
dis.readFully(buffer);
- var byteBuffer = ByteBuffer.wrap(buffer).order(ByteOrder.LITTLE_ENDIAN);
- var vector = new float[dimension];
- var floatBuffer = byteBuffer.asFloatBuffer();
- floatBuffer.get(vector);
- vectors.add(vectorTypeSupport.createFloatVector(vector));
+ ByteBuffer byteBuffer =
+ ByteBuffer.wrap(buffer).order(ByteOrder.LITTLE_ENDIAN);
+
+ float[] data = new float[dimension];
+ byteBuffer.asFloatBuffer().get(data);
+
+ vectors.add(vectorTypeSupport.createFloatVector(data));
}
}
return vectors;
}
+ /**
+ * Read at most {@code numVectors} vectors from an .fvecs file.
+ *
+ *
+ * This method stops reading once {@code numVectors} vectors
+ * have been loaded, preventing unnecessary heap usage.
+ *
+ *
+ *
+ * If the file contains fewer than {@code numVectors} vectors, an exception is thrown.
+ *
+ *
+ * @param filePath path to the .fvecs file
+ * @param numVectors number of vectors to load
+ * @return list of vectors of size {@code numVectors}
+ *
+ * @throws IllegalArgumentException if {@code numVectors <= 0}
+ * @throws IOException if the file ends before {@code numVectors}
+ * vectors are read
+ */
+ public static List> readFvecs(String filePath, int numVectors)
+ throws IOException {
+
+ if (numVectors <= 0) {
+ throw new IllegalArgumentException(
+ "numVectors must be > 0, got " + numVectors);
+ }
+
+ var vectors = new ArrayList>(
+ Math.min(numVectors, 1024));
+
+ try (var dis = new DataInputStream(
+ new BufferedInputStream(new FileInputStream(filePath)))) {
+
+ int count = 0;
+
+ while (count < numVectors) {
+ if (dis.available() <= 0) {
+ throw new IOException(
+ "File ended early while reading " + filePath +
+ ": requested " + numVectors +
+ " vectors, but only read " + count);
+ }
+
+ int dimension = Integer.reverseBytes(dis.readInt());
+ if (dimension <= 0) {
+ throw new IOException("Invalid vector dimension: " + dimension);
+ }
+
+ byte[] buffer = new byte[dimension * Float.BYTES];
+ dis.readFully(buffer);
+
+ ByteBuffer byteBuffer =
+ ByteBuffer.wrap(buffer).order(ByteOrder.LITTLE_ENDIAN);
+
+ float[] data = new float[dimension];
+ byteBuffer.asFloatBuffer().get(data);
+
+ vectors.add(vectorTypeSupport.createFloatVector(data));
+ count++;
+ }
+ }
+
+ // Defensive sanity check
+ if (vectors.size() != numVectors) {
+ throw new IOException(
+ "Internal error: requested " + numVectors +
+ " vectors but read " + vectors.size());
+ }
+
+ return vectors;
+ }
+
+ /**
+ * Read all ground-truth neighbors from an .ivecs file.
+ */
public static List> readIvecs(String filename) {
var groundTruthTopK = new ArrayList>();
try (var dis = new DataInputStream(new FileInputStream(filename))) {
while (dis.available() > 0) {
- var numNeighbors = Integer.reverseBytes(dis.readInt());
+ int numNeighbors = Integer.reverseBytes(dis.readInt());
var neighbors = new ArrayList(numNeighbors);
- for (var i = 0; i < numNeighbors; i++) {
- var neighbor = Integer.reverseBytes(dis.readInt());
- neighbors.add(neighbor);
+ for (int i = 0; i < numNeighbors; i++) {
+ neighbors.add(Integer.reverseBytes(dis.readInt()));
}
groundTruthTopK.add(neighbors);
@@ -74,4 +168,41 @@ public static List> readIvecs(String filename) {
return groundTruthTopK;
}
+
+ // Fast, allocation-free scan to count vectors in an .fvecs file
+ public static int countFvecs(String filePath) throws IOException {
+ Path path = Path.of(filePath);
+
+ long fileSize = Files.size(path);
+ if (fileSize < Integer.BYTES) {
+ throw new IOException("File too small to be a valid .fvecs file: " + filePath);
+ }
+
+ // Read dimension from first record
+ int dimension;
+ try (FileChannel fc = FileChannel.open(path, StandardOpenOption.READ)) {
+ ByteBuffer buf = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN);
+ fc.read(buf);
+ buf.flip();
+ dimension = buf.getInt();
+ }
+
+ if (dimension <= 0) {
+ throw new IOException("Invalid vector dimension: " + dimension);
+ }
+
+ long recordBytes = Integer.BYTES + (long) dimension * Float.BYTES;
+ if (fileSize % recordBytes != 0) {
+ throw new IOException(
+ "File size not a multiple of record size: " +
+ "fileSize=" + fileSize + ", recordBytes=" + recordBytes);
+ }
+
+ long count = fileSize / recordBytes;
+ if (count > Integer.MAX_VALUE) {
+ throw new IOException("Vector count exceeds Integer.MAX_VALUE: " + count);
+ }
+
+ return (int) count;
+ }
}
From d8aa4f663fa726a7b7dc513f144d23363d079709 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Fri, 2 Jan 2026 16:46:29 -0800
Subject: [PATCH 013/243] Added tput reporting to DistancesASH and DistancesPQ.
---
.../jbellis/jvector/example/DistancesASH.java | 39 +++++++++++++++----
.../jbellis/jvector/example/DistancesPQ.java | 27 ++++++++++++-
2 files changed, 56 insertions(+), 10 deletions(-)
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index ac8fe886a..d61d904c5 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -61,11 +61,11 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {32}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {16, 32, 64}; // 16, 32, and/or 64
// Define the benchmark size
- int maxQueries = 10_000;
- int maxVectors = 3_000_000;
+ int maxQueries = 1_000;
+ int maxVectors = 100_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -143,6 +143,15 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
long endTime = System.nanoTime();
System.out.println("\tEncoding took " + (endTime - startTime) / 1e9 + " seconds");
+ double encSeconds = (endTime - startTime) / 1e9;
+ double encThroughput = finalVectorCount / encSeconds;
+
+ System.out.println(
+ "\tEncoding throughput = "
+ + String.format(java.util.Locale.ROOT, "%.3f", encThroughput)
+ + " vectors/sec"
+ );
+
ASHVectors ashVectors = (ASHVectors) ashVecs;
VectorFloat> q0 = finalQueries.get(0);
@@ -389,6 +398,20 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
+ (blockEnd - blockStart) / 1e9 + " seconds"
);
+ double blockSeconds = (blockEnd - blockStart) / 1e9;
+ long totalDotProducts = (long) finalQueryCount * (long) finalVectorCount;
+
+ double scoreThroughput = totalDotProducts / blockSeconds;
+
+ System.out.println(
+ "\tBlock " + kernelMode + " throughput (blockSize=" + blockSize + ") = "
+ + String.format(java.util.Locale.ROOT, "%.3f", scoreThroughput)
+ + " dot-products/sec"
+ + " ("
+ + String.format(java.util.Locale.ROOT, "%.3f", scoreThroughput / 1e6)
+ + " Mdot/s)"
+ );
+
System.out.println("\tdummyAccumulator = " + (float) (blockDummy));
System.out.println("--");
}
@@ -470,11 +493,11 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
-// runCohere100k();
-// runADA();
-// runOpenai1536();
-// runOpenai3072();
+ runCohere100k();
+ runADA();
+ runOpenai1536();
+ runOpenai3072();
// runCap6m();
- runCohere10m();
+// runCohere10m();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
index a11f3c156..6c5fe9190 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -57,8 +57,8 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
// Define the benchmark size
- int maxQueries = 10_000;
- int maxVectors = 2_000_000;
+ int maxQueries = 1_000;
+ int maxVectors = 100_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -162,6 +162,15 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
System.out.println("\tPQ training took " + (pqTrainEnd - pqTrainStart) / 1e9 + " seconds");
System.out.println("\tPQ encoding took " + (pqEncodeEnd - pqEncodeStart) / 1e9 + " seconds");
+ double encSeconds = (pqEncodeEnd - pqEncodeStart) / 1e9;
+ double encThroughput = finalVectorCount / encSeconds;
+
+ System.out.println(
+ "\tEncoding throughput = "
+ + String.format(java.util.Locale.ROOT, "%.3f", encThroughput)
+ + " vectors/sec"
+ );
+
// ============================================================
// [1] Accuracy (NOT timed)
@@ -245,6 +254,20 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
long pqEnd = System.nanoTime();
System.out.println("\tPQDecoder scan took " + (pqEnd - pqStart) / 1e9 + " seconds");
+ double blockSeconds = (pqEnd - pqStart) / 1e9;
+ long totalDotProducts = (long) finalQueryCount * (long) finalVectorCount;
+
+ double scoreThroughput = totalDotProducts / blockSeconds;
+
+ System.out.println(
+ "\tDecoder scan throughput = "
+ + String.format(java.util.Locale.ROOT, "%.3f", scoreThroughput)
+ + " dot-products/sec"
+ + " ("
+ + String.format(java.util.Locale.ROOT, "%.3f", scoreThroughput / 1e6)
+ + " Mdot/s)"
+ );
+
// ============================================================
// [3] Float dot-product baseline
// ============================================================
From 3e5b4e39f31d83a263b33a77184d65c848016ca7 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Fri, 2 Jan 2026 16:47:49 -0800
Subject: [PATCH 014/243] First phase of ASH multi-landmark support.
---
.../jvector/quantization/ASHScorer.java | 3 +-
.../jvector/quantization/ASHVectors.java | 4 +-
.../quantization/AsymmetricHashing.java | 117 ++++++++++++------
3 files changed, 85 insertions(+), 39 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
index f3a2cfea9..fc6925eb1 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
@@ -51,7 +51,7 @@ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
// Compute q_centered = q - μ
// (avoid allocation by doing subtraction on the fly during projection)
- final VectorFloat> mu = ash.globalMean;
+ final VectorFloat> mu = ash.landmarks[0];
final float[][] A = ash.stiefelTransform.AFloat; // shape: d × originalDim
final int originalDim = ash.stiefelTransform.cols;
@@ -138,4 +138,3 @@ public interface ASHScoreFunction {
float similarityTo(AsymmetricHashing.QuantizedVector vector2);
}
}
-
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index 86277d4d2..7d92d3635 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -380,7 +380,7 @@ private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
final int D = ash.originalDimension;
- final VectorFloat> mu = ash.globalMean;
+ final VectorFloat> mu = ash.landmarks[0];
final float[] qArr = new float[D];
final float[] muArr = new float[D];
@@ -513,7 +513,7 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
// ---- Query precompute (identical to ASHScorer) ----
final int D = ash.originalDimension;
- final VectorFloat> mu = ash.globalMean;
+ final VectorFloat> mu = ash.landmarks[0];
final float[] qArr = new float[D];
final float[] muArr = new float[D];
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 9f5ca260a..e43d0f82d 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -34,8 +34,6 @@
import java.util.Objects;
import java.util.concurrent.ForkJoinPool;
import java.util.Random;
-import java.util.stream.IntStream;
-import org.apache.commons.math3.linear.Array2DRowRealMatrix;
import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.QRDecomposition;
import org.apache.commons.math3.linear.RealMatrix;
@@ -65,8 +63,11 @@ public class AsymmetricHashing implements VectorCompressor globalMean;
+ /** Number of landmarks (clusters), C<= 64 */
+ public final int landmarkCount;
+
+ /** Landmark centroids μ_0... μ_{C−1} */
+ public final VectorFloat>[] landmarks;
/** Original (uncompressed) dimensionality */
public final int originalDimension;
@@ -111,25 +112,35 @@ private BinaryQuantizer quantizer() {
/**
* Class constructor.
* @param originalDim the dimensionality of the vectors to be quantized
- * @param globalMean the mean of the database (its average vector)
+ * @param landmarks the landmarks using for vectors (up to C in landmarkCount)
*/
private AsymmetricHashing(int originalDim,
int encodedBits,
int quantizedDim,
int optimizer,
- VectorFloat> globalMean,
+ VectorFloat>[] landmarks,
StiefelTransform stiefelTransform) {
this.originalDimension = originalDim;
this.encodedBits = encodedBits;
this.quantizedDim = quantizedDim;
this.optimizer = optimizer;
- this.stiefelTransform = Objects.requireNonNull(stiefelTransform, "stiefelTransform");
- this.globalMean = Objects.requireNonNull(globalMean, "globalMean");
+ this.stiefelTransform = Objects.requireNonNull(stiefelTransform);
+
+ this.landmarks = Objects.requireNonNull(landmarks);
+ this.landmarkCount = landmarks.length;
- if (globalMean.length() != originalDimension) {
+ if (landmarkCount < 1 || landmarkCount > 64) {
throw new IllegalArgumentException(
- "Global mean length does not match original dimension");
+ "landmarkCount must be in [1,64], got " + landmarkCount);
+ }
+
+ for (VectorFloat> mu : landmarks) {
+ if (mu.length() != originalDimension) {
+ throw new IllegalArgumentException(
+ "Landmark dimensionality does not match original dimensionality");
+ }
}
+
if (quantizedDim < 1) {
throw new IllegalArgumentException("quantizedDim must be > 0");
}
@@ -159,12 +170,14 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
var ravvCopy = ravv.threadLocalSupplier().get();
int originalDim = ravvCopy.getVector(0).length();
- // Compute global mean
- var globalMean = vectorTypeSupport.createFloatVector(originalDim);
+ // C=1 landmarks case
+ VectorFloat> mu0 = vectorTypeSupport.createFloatVector(originalDim);
for (int i = 0; i < ravvCopy.size(); i++) {
- VectorUtil.addInPlace(globalMean, ravvCopy.getVector(i));
+ VectorUtil.addInPlace(mu0, ravvCopy.getVector(i));
}
- VectorUtil.scale(globalMean, 1.0f / ravvCopy.size());
+ VectorUtil.scale(mu0, 1.0f / ravvCopy.size());
+
+ VectorFloat>[] landmarks = new VectorFloat>[] { mu0 };
StiefelTransform stiefelTransform;
if (optimizer == RANDOM) {
@@ -182,7 +195,7 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
encodedBits,
quantizedDim,
optimizer,
- globalMean,
+ landmarks,
stiefelTransform
);
}
@@ -210,11 +223,12 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
int quantizedDim = in.readInt();
int optimizer = in.readInt();
- int globalMeanLength = in.readInt();
- if (globalMeanLength <= 0) {
- throw new IOException("ASH header missing globalMean");
+ int landmark0Length = in.readInt();
+ if (landmark0Length <= 0) {
+ throw new IOException("ASH header missing landmark[0]");
}
- VectorFloat> globalMean = vectorTypeSupport.readFloatVector(in, globalMeanLength);
+ VectorFloat> mu0 = vectorTypeSupport.readFloatVector(in, landmark0Length);
+ VectorFloat>[] landmarks = new VectorFloat>[] { mu0 };
// NOTE: StiefelTransform loading will be added in Step 2.
// For now, construct a placeholder or re-run the RANDOM initializer.
@@ -238,7 +252,7 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
encodedBits,
quantizedDim,
optimizer,
- globalMean,
+ landmarks,
stiefelTransform
);
}
@@ -264,8 +278,9 @@ public void write(DataOutput out, int version) throws IOException {
out.writeInt(quantizedDim);
out.writeInt(optimizer);
- out.writeInt(globalMean.length());
- vectorTypeSupport.writeFloatVector(out, globalMean);
+ // Step 1: write a single landmark (mu0) in the same slot as legacy globalMean
+ out.writeInt(landmarks[0].length());
+ vectorTypeSupport.writeFloatVector(out, landmarks[0]);
// NOTE: stiefelTransform serialization will be added later
}
@@ -286,13 +301,20 @@ public void quantizeVector(VectorFloat> vector,
") exceeds original dimension (" + originalDimension + ")");
}
+ if ((landmark & 0xFFFF) >= landmarkCount) {
+ throw new IllegalArgumentException("Invalid landmark id " + (landmark & 0xFFFF)
+ + " for landmarkCount=" + landmarkCount);
+ }
+
+ final VectorFloat> mu = landmarks[landmark]; // Step 1 landmark==0 always, but keep correct
+
QuantizedVector.quantizeTo(
vector,
residualNorm,
dotWithLandmark,
landmark,
quantizedDim,
- globalMean,
+ mu,
stiefelTransform,
quantizer(),
dest
@@ -806,14 +828,17 @@ public void encodeTo(VectorFloat> vector, QuantizedVector dest) {
// All vectors are encoded relative to the global mean μ*.
// The landmark id is always 0 in this mode.
- final short landmark = 0;
+ // Landmark assignment (C=1 → always 0)
+ final short landmark = 0; // placeholder; Step 2 will generalize
+
+ final VectorFloat> mu = landmarks[landmark];
- // Table 2 header field:
- final float dotWithLandmark = VectorUtil.dotProduct(vector, globalMean);
+ // ASH paper, Table 2 header field:
+ final float dotWithLandmark = VectorUtil.dotProduct(vector, mu);
// Table 2 header field: ||x_i − μ_i*||₂
// We store the true L2 norm (not squared), matching the paper.
- final float sqDist = VectorUtil.squareL2Distance(vector, globalMean);
+ final float sqDist = VectorUtil.squareL2Distance(vector, mu);
final float residualNorm = (float) Math.sqrt(sqDist);
// Sanity: quantizedDim already accounts for physical header bits (Step 1)
@@ -834,8 +859,8 @@ public int compressorSize() {
size += Integer.BYTES; // encodedBits
size += Integer.BYTES; // quantizedDim
size += Integer.BYTES; // optimizer
- size += Integer.BYTES; // globalMean length
- size += Float.BYTES * globalMean.length();
+ size += Integer.BYTES; // landmarks[0] length
+ size += Float.BYTES * landmarks[0].length();
return size;
}
@@ -857,8 +882,10 @@ public double reconstructionError(VectorFloat> vector) {
public long ramBytesUsed() {
long bytes = 0L;
- if (globalMean != null) {
- bytes += globalMean.ramBytesUsed();
+ if (landmarks != null) {
+ for (int i = 0; i < landmarks.length; i++) {
+ bytes += (landmarks[i] == null) ? 0L : landmarks[i].ramBytesUsed();
+ }
}
if (stiefelTransform != null) {
@@ -1001,13 +1028,23 @@ public boolean equals(Object o) {
&& encodedBits == that.encodedBits
&& quantizedDim == that.quantizedDim
&& optimizer == that.optimizer
- && vectorEquals(this.globalMean, that.globalMean)
+ && landmarksEqual(this.landmarks, that.landmarks)
&& matrixEquals(
this.stiefelTransform == null ? null : this.stiefelTransform.W,
that.stiefelTransform == null ? null : that.stiefelTransform.W
);
}
+ private static boolean landmarksEqual(VectorFloat>[] a, VectorFloat>[] b) {
+ if (a == b) return true;
+ if (a == null || b == null) return false;
+ if (a.length != b.length) return false;
+ for (int i = 0; i < a.length; i++) {
+ if (!vectorEquals(a[i], b[i])) return false;
+ }
+ return true;
+ }
+
/**
* Hash code consistent with {@link #equals(Object)}.
*
@@ -1023,13 +1060,22 @@ public int hashCode() {
result = 31 * result + Integer.hashCode(encodedBits);
result = 31 * result + Integer.hashCode(quantizedDim);
result = 31 * result + Integer.hashCode(optimizer);
- result = 31 * result + vectorHash(globalMean);
+ result = 31 * result + landmarksHash(landmarks);
result = 31 * result + matrixHash(
stiefelTransform == null ? null : stiefelTransform.W
);
return result;
}
+ private static int landmarksHash(VectorFloat>[] a) {
+ if (a == null) return 0;
+ int h = 1;
+ for (VectorFloat> v : a) {
+ h = 31 * h + vectorHash(v);
+ }
+ return h;
+ }
+
@Override
public String toString() {
String optimizerName;
@@ -1048,11 +1094,12 @@ public String toString() {
}
return String.format(
- "AsymmetricHashing[origDim=%d, encodedBits=%d, quantizedDim=%d, optimizer=%s, landmarks=1]",
+ "AsymmetricHashing[origDim=%d, encodedBits=%d, quantizedDim=%d, optimizer=%s, landmarks=%d]",
originalDimension,
encodedBits,
quantizedDim,
- optimizerName
+ optimizerName,
+ landmarkCount
);
}
}
From 3869e3644a1c988f31ba24578022b71170e7231b Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Sat, 3 Jan 2026 06:26:18 +0000
Subject: [PATCH 015/243] More multi-landmark refactoring and moved the full
scale and offset into the vector header.
---
.../jvector/quantization/ASHScorer.java | 15 +-
.../jvector/quantization/ASHVectors.java | 66 ++++----
.../quantization/AsymmetricHashing.java | 155 +++++++++++-------
3 files changed, 133 insertions(+), 103 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
index fc6925eb1..2bd37c4e1 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
@@ -88,19 +88,10 @@ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
//
final float dotQMu = VectorUtil.dotProduct(query, mu);
- // ||μ||^2
- final float muNormSq = VectorUtil.dotProduct(mu, mu);
-
- // Precompute invSqrtD once
- final float invSqrtD = (float) (1.0 / Math.sqrt(d));
-
return (AsymmetricHashing.QuantizedVector v) -> {
// Single landmark baseline: ignore v.landmark for now (assumed 0)
- final float residualNorm = v.residualNorm; // ||x - μ||₂
- final float dotXMu = v.dotWithLandmark; //
-
- // scale = d^{-1/2} * ||x - μ||
- final float scale = invSqrtD * residualNorm;
+ final float scale = v.scale; // already ||x − μ|| / sqrt(d)
+ final float offset = v.offset; // already − ||μ||^2
// masked-add term: where bin ∈ {0,1}^d
// This is sum of tildeQ[j] over set bits in v.binaryVector.
@@ -128,7 +119,7 @@ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
// - scale = ||x − μ|| / √d (pure Eq. 11 baseline; no bias correction).
// - ⟨q, μ⟩ is a query-only offset and MUST NOT be scaled.
// - Scaling ⟨q, μ⟩ leads to large systematic error (see paper Eq. 11 regrouping).
- return scale * (2f * maskedAdd - sumTildeQ) + dotQMu + (dotXMu - muNormSq);
+ return scale * (2f * maskedAdd - sumTildeQ) + dotQMu + offset;
};
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index 7d92d3635..b4c1e6bcd 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -71,8 +71,8 @@ static AshBlockKernel fromProperty() {
}
// Cached scalar headers (avoid object chasing in blocked scorer)
- private final float[] residualNorms;
- private final float[] dotXMu;
+ private final float[] scales; // scale_i = ||x_i − μ|| / sqrt(d)
+ private final float[] offsets; // offset_i = − ||μ||^2
// Quantized dimensions
private final int d;
@@ -99,12 +99,12 @@ public ASHVectors(AsymmetricHashing ash,
this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
int n = compressedVectors.length;
- this.residualNorms = new float[n];
- this.dotXMu = new float[n];
+ this.scales = new float[n];
+ this.offsets = new float[n];
for (int i = 0; i < n; i++) {
var v = compressedVectors[i];
- residualNorms[i] = v.residualNorm;
- dotXMu[i] = v.dotWithLandmark;
+ scales[i] = v.scale;
+ offsets[i] = v.offset;
}
}
@@ -285,8 +285,8 @@ public ASHBlockScorer blockScorerFor(
return new BlockScalarASHBlockScorer(
ash,
compressedVectors,
- residualNorms,
- dotXMu,
+ scales,
+ offsets,
query
);
}
@@ -302,8 +302,8 @@ public ASHBlockScorer blockScorerFor(
return new BlockSimdASHBlockScorer(
ash,
compressedVectors,
- residualNorms,
- dotXMu,
+ scales,
+ offsets,
packedBits,
blockSize,
query
@@ -352,14 +352,13 @@ private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
private final AsymmetricHashing ash;
private final AsymmetricHashing.QuantizedVector[] vectors;
- private final float[] residualNorms;
- private final float[] dotXMu;
+ private final float[] scales;
+ private final float[] offsets;
+
private final float[] tildeQ;
private final float sumTildeQ;
private final float dotQMu;
- private final float muNormSq;
- private final float invSqrtD;
private final int d;
private final int words;
@@ -367,14 +366,14 @@ private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
BlockScalarASHBlockScorer(
AsymmetricHashing ash,
AsymmetricHashing.QuantizedVector[] vectors,
- float[] residualNorms,
- float[] dotXMu,
+ float[] scales,
+ float[] offsets,
VectorFloat> query) {
this.ash = ash;
this.vectors = vectors;
- this.residualNorms = residualNorms;
- this.dotXMu = dotXMu;
+ this.scales = scales;
+ this.offsets = offsets;
this.d = ash.quantizedDim;
this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
@@ -403,10 +402,7 @@ private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
float sum = 0f;
for (int i = 0; i < d; i++) sum += tildeQ[i];
this.sumTildeQ = sum;
-
this.dotQMu = VectorUtil.dotProduct(query, mu);
- this.muNormSq = VectorUtil.dotProduct(mu, mu);
- this.invSqrtD = (float) (1.0 / Math.sqrt(d));
}
@Override
@@ -430,11 +426,13 @@ public void scoreRange(int start, int count, float[] out) {
}
}
- float scale = invSqrtD * residualNorms[idx];
+ float scale = scales[idx];
+ float offset = offsets[idx];
+
out[i] =
scale * (2f * maskedAdd - sumTildeQ)
+ dotQMu
- + (dotXMu[idx] - muNormSq);
+ + offset;
}
}
}
@@ -465,8 +463,8 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private final AsymmetricHashing ash;
private final AsymmetricHashing.QuantizedVector[] vectors;
- private final float[] residualNorms;
- private final float[] dotXMu;
+ private final float[] scales;
+ private final float[] offsets;
private final long[] packedBits;
private final int blockSize;
@@ -478,8 +476,6 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private final float[] tildeQ;
private final float sumTildeQ;
private final float dotQMu;
- private final float muNormSq;
- private final float invSqrtD;
// Scratch (one per scorer instance)
private final float[] maskedAdds;
@@ -489,8 +485,8 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
BlockSimdASHBlockScorer(
AsymmetricHashing ash,
AsymmetricHashing.QuantizedVector[] vectors,
- float[] residualNorms,
- float[] dotXMu,
+ float[] scales,
+ float[] offsets,
long[] packedBits,
int blockSize,
VectorFloat> query) {
@@ -501,8 +497,8 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
.getVectorUtilSupport();
this.ash = ash;
this.vectors = vectors;
- this.residualNorms = residualNorms;
- this.dotXMu = dotXMu;
+ this.scales = scales;
+ this.offsets = offsets;
this.packedBits = packedBits;
this.blockSize = blockSize;
@@ -536,10 +532,7 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
float sum = 0f;
for (int i = 0; i < d; i++) sum += tildeQ[i];
this.sumTildeQ = sum;
-
this.dotQMu = VectorUtil.dotProduct(query, mu);
- this.muNormSq = VectorUtil.dotProduct(mu, mu);
- this.invSqrtD = (float) (1.0 / Math.sqrt(d));
}
@Override
@@ -573,13 +566,14 @@ public void scoreRange(int start, int count, float[] out) {
// Finish Eq. 11 per vector
for (int lane = 0; lane < blockLen; lane++) {
int idx = ord + lane;
- float scale = invSqrtD * residualNorms[idx];
+ float scale = scales[idx];
+ float offset = offsets[idx];
float m = maskedAdds[lane];
out[idx - start] =
scale * (2f * m - sumTildeQ)
+ dotQMu
- + (dotXMu[idx] - muNormSq);
+ + offset;
}
ord += blockLen;
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index e43d0f82d..4fe0383f0 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -50,9 +50,12 @@ public class AsymmetricHashing implements VectorCompressor − ||μ||_2^2
+ // - landmark id: short (16 bits) in [0, C)
+ //
+ // NOTE: residualNorm (= ||x − μ||) is still computed during encoding for Eq. 6,
+ // but it is not stored explicitly once scale is stored.
private static final int HEADER_BITS =
(Float.BYTES + Float.BYTES + Short.BYTES) * 8; // 80 bits currently
@@ -69,6 +72,9 @@ public class AsymmetricHashing implements VectorCompressor[] landmarks;
+ // One entry per landmark
+ private final float[] landmarkNormSq;
+
/** Original (uncompressed) dimensionality */
public final int originalDimension;
@@ -127,8 +133,20 @@ private AsymmetricHashing(int originalDim,
this.stiefelTransform = Objects.requireNonNull(stiefelTransform);
this.landmarks = Objects.requireNonNull(landmarks);
+ assert landmarks.length > 0;
this.landmarkCount = landmarks.length;
+ // Precompute ||μ_c||^2 for each landmark (even when C = 1)
+ this.landmarkNormSq = new float[landmarkCount];
+ for (int c = 0; c < landmarkCount; c++) {
+ this.landmarkNormSq[c] =
+ VectorUtil.dotProduct(landmarks[c], landmarks[c]);
+ }
+
+ if (landmarkNormSq.length != landmarkCount) {
+ throw new IllegalArgumentException("landmarkNormSq length mismatch");
+ }
+
if (landmarkCount < 1 || landmarkCount > 64) {
throw new IllegalArgumentException(
"landmarkCount must be in [1,64], got " + landmarkCount);
@@ -278,7 +296,7 @@ public void write(DataOutput out, int version) throws IOException {
out.writeInt(quantizedDim);
out.writeInt(optimizer);
- // Step 1: write a single landmark (mu0) in the same slot as legacy globalMean
+ // (C=1): write landmark[0] (dataset mean). Multi-landmark serialization TBD.
out.writeInt(landmarks[0].length());
vectorTypeSupport.writeFloatVector(out, landmarks[0]);
@@ -291,7 +309,6 @@ public void write(DataOutput out, int version) throws IOException {
public void quantizeVector(VectorFloat> vector,
float residualNorm,
- float dotWithLandmark,
short landmark,
QuantizedVector dest) {
@@ -306,13 +323,12 @@ public void quantizeVector(VectorFloat> vector,
+ " for landmarkCount=" + landmarkCount);
}
- final VectorFloat> mu = landmarks[landmark]; // Step 1 landmark==0 always, but keep correct
+ // Landmark-specific mean (C=1 → landmarks[0] == dataset mean)
+ final VectorFloat> mu = landmarks[landmark];
QuantizedVector.quantizeTo(
vector,
residualNorm,
- dotWithLandmark,
- landmark,
quantizedDim,
mu,
stiefelTransform,
@@ -327,7 +343,7 @@ public void quantizeVector(VectorFloat> vector,
public interface BinaryQuantizer {
void quantizeBody(VectorFloat> vector,
- VectorFloat> globalMean,
+ VectorFloat> mu,
float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
@@ -338,7 +354,7 @@ public static final class RandomBinaryQuantizer implements BinaryQuantizer {
@Override
public void quantizeBody(VectorFloat> vector,
- VectorFloat> globalMean,
+ VectorFloat> mu,
float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
@@ -354,10 +370,10 @@ public void quantizeBody(VectorFloat> vector,
// Copy vector and mean once (no virtual calls in inner loop)
final float[] x = new float[originalDim];
- final float[] mu = new float[originalDim];
+ final float[] muArr = new float[originalDim];
for (int d = 0; d < originalDim; d++) {
x[d] = vector.get(d);
- mu[d] = globalMean.get(d);
+ muArr[d] = mu.get(d);
}
// ASH paper, Eq. 6, normalization
@@ -366,7 +382,7 @@ public void quantizeBody(VectorFloat> vector,
// Normalized residual x̂
final float[] xhat = new float[originalDim];
for (int d = 0; d < originalDim; d++) {
- xhat[d] = (x[d] - mu[d]) * invNorm;
+ xhat[d] = (x[d] - muArr[d]) * invNorm;
}
// Binarize directly from per-row projection
@@ -391,7 +407,7 @@ public void quantizeBody(VectorFloat> vector,
public static final class ItqBinaryQuantizer implements BinaryQuantizer {
@Override
public void quantizeBody(VectorFloat> vector,
- VectorFloat> globalMean,
+ VectorFloat> mu,
float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
@@ -403,7 +419,7 @@ public void quantizeBody(VectorFloat> vector,
public static final class LandingBinaryQuantizer implements BinaryQuantizer {
@Override
public void quantizeBody(VectorFloat> vector,
- VectorFloat> globalMean,
+ VectorFloat> mu,
float residualNorm,
int quantizedDim,
StiefelTransform stiefel,
@@ -417,17 +433,17 @@ public void quantizeBody(VectorFloat> vector,
// ---------------------------------------------------------------------
public static class QuantizedVector {
- public float residualNorm; // ||x_i - μ_i*||_2
- public float dotWithLandmark; //
+ public float scale; // d^{-1/2} ||x_i - μ_i*||_2
+ public float offset; // offset_i = - ||μ_i*||^2
public short landmark; // c_i*
public long[] binaryVector; // bin(x̂_i)
- private QuantizedVector(float residualNorm,
- float dotWithLandmark,
+ private QuantizedVector(float scale,
+ float offset,
short landmark,
long[] binaryVector) {
- this.residualNorm = residualNorm;
- this.dotWithLandmark = dotWithLandmark;
+ this.scale = scale;
+ this.offset = offset;
this.landmark = landmark;
this.binaryVector = binaryVector;
}
@@ -438,8 +454,11 @@ public static int wordsForDims(int quantizedDim) {
public static QuantizedVector createEmpty(int quantizedDim) {
return new QuantizedVector(
- 0f, 0f, (short) 0,
- new long[wordsForDims(quantizedDim)]);
+ Float.NaN,
+ Float.NaN,
+ (short) 0,
+ new long[wordsForDims(quantizedDim)]
+ );
}
public static int serializedSizeBytes(int quantizedDim) {
@@ -447,19 +466,31 @@ public static int serializedSizeBytes(int quantizedDim) {
+ wordsForDims(quantizedDim) * Long.BYTES;
}
+ /**
+ * Writes only the binary body into {@code dest.binaryVector}.
+ *
+ * Header fields ({@code scale}, {@code offset}, {@code landmark}) must be
+ * set by the caller prior to calling this method.
+ *
+ * @param residualNorm ||x - μ||_2, required for Eq. 6 normalization during binarization
+ */
static void quantizeTo(VectorFloat> vector,
float residualNorm,
- float dotWithLandmark,
- short landmark,
int quantizedDim,
- VectorFloat> globalMean,
+ VectorFloat> mu,
StiefelTransform stiefel,
BinaryQuantizer quantizer,
QuantizedVector dest) {
- dest.residualNorm = residualNorm;
- dest.dotWithLandmark = dotWithLandmark;
- dest.landmark = landmark;
+ // encodeTo() must have initialized header
+ if (Float.isNaN(dest.scale) || Float.isNaN(dest.offset)) {
+ throw new IllegalStateException(
+ "QuantizedVector header not initialized before quantizeTo()"
+ );
+ }
+
+ assert Float.isFinite(dest.scale) : "scale is not finite";
+ assert Float.isFinite(dest.offset) : "offset is not finite";
int words = wordsForDims(quantizedDim);
if (dest.binaryVector.length < words) {
@@ -469,7 +500,7 @@ static void quantizeTo(VectorFloat> vector,
Arrays.fill(dest.binaryVector, 0, words, 0L);
quantizer.quantizeBody(
vector,
- globalMean,
+ mu,
residualNorm,
quantizedDim,
stiefel,
@@ -478,8 +509,8 @@ static void quantizeTo(VectorFloat> vector,
}
public void write(DataOutput out, int quantizedDim) throws IOException {
- out.writeFloat(residualNorm);
- out.writeFloat(dotWithLandmark);
+ out.writeFloat(scale);
+ out.writeFloat(offset);
out.writeInt(landmark);
int words = wordsForDims(quantizedDim);
@@ -490,8 +521,8 @@ public void write(DataOutput out, int quantizedDim) throws IOException {
public static QuantizedVector load(RandomAccessReader in,
int quantizedDim) throws IOException {
- float residualNorm = in.readFloat();
- float dotWithLandmark = in.readFloat();
+ float scale = in.readFloat();
+ float offset = in.readFloat();
short landmark = (short) in.readInt();
int words = wordsForDims(quantizedDim);
@@ -500,14 +531,14 @@ public static QuantizedVector load(RandomAccessReader in,
body[i] = in.readLong();
}
- return new QuantizedVector(residualNorm, dotWithLandmark, landmark, body);
+ return new QuantizedVector(scale, offset, landmark, body);
}
public static void loadInto(RandomAccessReader in,
QuantizedVector dest,
int quantizedDim) throws IOException {
- dest.residualNorm = in.readFloat();
- dest.dotWithLandmark = in.readFloat();
+ dest.scale = in.readFloat();
+ dest.offset = in.readFloat();
dest.landmark = (short) in.readInt();
int words = wordsForDims(quantizedDim);
@@ -540,8 +571,8 @@ public boolean equals(Object o) {
if (o == null || getClass() != o.getClass()) return false;
QuantizedVector that = (QuantizedVector) o;
- return Float.floatToIntBits(residualNorm) == Float.floatToIntBits(that.residualNorm)
- && Float.floatToIntBits(dotWithLandmark) == Float.floatToIntBits(that.dotWithLandmark)
+ return Float.floatToIntBits(scale) == Float.floatToIntBits(that.scale)
+ && Float.floatToIntBits(offset) == Float.floatToIntBits(that.offset)
&& landmark == that.landmark
&& Arrays.equals(binaryVector, that.binaryVector);
}
@@ -556,8 +587,8 @@ public boolean equals(Object o) {
*/
@Override
public int hashCode() {
- int result = Integer.hashCode(Float.floatToIntBits(residualNorm));
- result = 31 * result + Integer.hashCode(Float.floatToIntBits(dotWithLandmark));
+ int result = Integer.hashCode(Float.floatToIntBits(scale));
+ result = 31 * result + Integer.hashCode(Float.floatToIntBits(offset));
result = 31 * result + Short.hashCode(landmark);
result = 31 * result + Arrays.hashCode(binaryVector);
return result;
@@ -825,28 +856,42 @@ public QuantizedVector encode(VectorFloat> vector) {
public void encodeTo(VectorFloat> vector, QuantizedVector dest) {
// NOTE:
// Single-landmark baseline configuration (Table 2, C = 1).
- // All vectors are encoded relative to the global mean μ*.
+ // All vectors are encoded relative to landmark μ_c (C=1 => μ_0 is the dataset mean).
// The landmark id is always 0 in this mode.
- // Landmark assignment (C=1 → always 0)
- final short landmark = 0; // placeholder; Step 2 will generalize
+ // Landmark assignment (C=1 --> always 0)
+ final short landmark = 0; // placeholder; will generalize
final VectorFloat> mu = landmarks[landmark];
- // ASH paper, Table 2 header field:
- final float dotWithLandmark = VectorUtil.dotProduct(vector, mu);
+ // Compute
+ final float dotXMu = VectorUtil.dotProduct(vector, mu);
- // Table 2 header field: ||x_i − μ_i*||₂
- // We store the true L2 norm (not squared), matching the paper.
+ // Compute ||x_i − μ_i*||_2
+ // We store the true L2 norm (not squared), matching the paper’s use in Eq. 6 normalization.
final float sqDist = VectorUtil.squareL2Distance(vector, mu);
final float residualNorm = (float) Math.sqrt(sqDist);
- // Sanity: quantizedDim already accounts for physical header bits (Step 1)
+ // Sanity: quantizedDim already accounts for physical header bits
// Binary body uses exactly quantizedDim bits.
assert quantizedDim > 0;
- // Binary body: sign(A · (x − μ))
- quantizeVector(vector, residualNorm, dotWithLandmark, landmark, dest);
+ // Store:
+ // scale = ||x − μ|| / sqrt(d)
+ // offset = − ||μ||^2
+ // NOTE: we still compute residualNorm (||x − μ||) even though we don't store it alone,
+ // because the binarizer needs it for Eq. 6 normalization.
+ final float invSqrtD = (float) (1.0 / Math.sqrt(quantizedDim));
+ final float scale = residualNorm * invSqrtD;
+ final float offset = dotXMu - landmarkNormSq[landmark];
+
+ // Write the new header fields into dest (requires QuantizedVector fields renamed, see below)
+ dest.scale = scale;
+ dest.offset = offset;
+ dest.landmark = landmark;
+
+ // Binary body: sign(A · (x − μ)) with Eq. 6 normalization inside quantizer
+ quantizeVector(vector, residualNorm, landmark, dest);
}
@Override
@@ -867,8 +912,8 @@ public int compressorSize() {
@Override
public int compressedVectorSize() {
int words = QuantizedVector.wordsForDims(quantizedDim);
- return Float.BYTES // residualNorm
- + Float.BYTES // dotWithLandmark
+ return Float.BYTES // scale
+ + Float.BYTES // offset
+ Integer.BYTES // landmark
+ words * Long.BYTES; // binary payload
}
@@ -1007,7 +1052,7 @@ private static int matrixHash(RealMatrix m) {
*
* - Matching dimensional parameters
* - Identical optimizer selection
- * - Bitwise-equal global mean vectors
+ * - Bitwise-equal landmark vectors
* - Bitwise-equal Stiefel transform matrices
*
*
@@ -1050,7 +1095,7 @@ private static boolean landmarksEqual(VectorFloat>[] a, VectorFloat>[] b) {
*
*
* The hash incorporates all parameters and learned state that affect encoding,
- * including the contents of the global mean and Stiefel transform. This ensures
+ * including the contents of landmarks and Stiefel transform. This ensures
* safe use of ASH compressors as map keys or cache entries.
*
*/
From 249456cc23b9f6a3553a7f64dad5280b74b6c1db Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Sat, 3 Jan 2026 22:08:58 +0000
Subject: [PATCH 016/243] Change landmark vector header field to byte.
---
.../jbellis/jvector/disk/ByteReadable.java | 30 +++++++++++++++++
.../jvector/quantization/ASHScorer.java | 1 +
.../jvector/quantization/ASHVectors.java | 8 +++--
.../quantization/AsymmetricHashing.java | 33 ++++++++++++-------
4 files changed, 58 insertions(+), 14 deletions(-)
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/disk/ByteReadable.java
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ByteReadable.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ByteReadable.java
new file mode 100644
index 000000000..f907dc45a
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ByteReadable.java
@@ -0,0 +1,30 @@
+/*
+ * 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.disk;
+
+import java.io.IOException;
+
+/**
+ * Optional extension for RandomAccessReader implementations that
+ * support byte-level reads.
+ *
+ * This is used by ASH only and MUST NOT be assumed by general readers.
+ */
+public interface ByteReadable {
+ byte readByte() throws IOException;
+}
+
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
index 2bd37c4e1..7ba87b284 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
@@ -51,6 +51,7 @@ private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
// Compute q_centered = q - μ
// (avoid allocation by doing subtraction on the fly during projection)
+ // Landmark-dependent query projection (C=1 => still index 0)
final VectorFloat> mu = ash.landmarks[0];
final float[][] A = ash.stiefelTransform.AFloat; // shape: d × originalDim
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index b4c1e6bcd..4e6bc43df 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -73,6 +73,7 @@ static AshBlockKernel fromProperty() {
// Cached scalar headers (avoid object chasing in blocked scorer)
private final float[] scales; // scale_i = ||x_i − μ|| / sqrt(d)
private final float[] offsets; // offset_i = − ||μ||^2
+ private final byte[] landmarks;
// Quantized dimensions
private final int d;
@@ -101,11 +102,14 @@ public ASHVectors(AsymmetricHashing ash,
int n = compressedVectors.length;
this.scales = new float[n];
this.offsets = new float[n];
+ this.landmarks = new byte[n];
for (int i = 0; i < n; i++) {
var v = compressedVectors[i];
- scales[i] = v.scale;
- offsets[i] = v.offset;
+ scales[i] = v.scale;
+ offsets[i] = v.offset;
+ landmarks[i] = v.landmark;
}
+
}
@Override
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 4fe0383f0..5514038f7 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -52,12 +52,12 @@ public class AsymmetricHashing implements VectorCompressor − ||μ||_2^2
- // - landmark id: short (16 bits) in [0, C)
+ // - landmark id: byte (8 bits) in [0, C)
//
// NOTE: residualNorm (= ||x − μ||) is still computed during encoding for Eq. 6,
// but it is not stored explicitly once scale is stored.
private static final int HEADER_BITS =
- (Float.BYTES + Float.BYTES + Short.BYTES) * 8; // 80 bits currently
+ (Float.BYTES + Float.BYTES + Byte.BYTES) * 8; // 72 bits currently
private static final VectorTypeSupport vectorTypeSupport =
VectorizationProvider.getInstance().getVectorTypeSupport();
@@ -318,8 +318,8 @@ public void quantizeVector(VectorFloat> vector,
") exceeds original dimension (" + originalDimension + ")");
}
- if ((landmark & 0xFFFF) >= landmarkCount) {
- throw new IllegalArgumentException("Invalid landmark id " + (landmark & 0xFFFF)
+ if ((landmark & 0xFF) >= landmarkCount) {
+ throw new IllegalArgumentException("Invalid landmark id " + (landmark & 0xFF)
+ " for landmarkCount=" + landmarkCount);
}
@@ -435,12 +435,12 @@ public void quantizeBody(VectorFloat> vector,
public static class QuantizedVector {
public float scale; // d^{-1/2} ||x_i - μ_i*||_2
public float offset; // offset_i = - ||μ_i*||^2
- public short landmark; // c_i*
+ public byte landmark; // c_i*, unsigned [0, C)
public long[] binaryVector; // bin(x̂_i)
private QuantizedVector(float scale,
float offset,
- short landmark,
+ byte landmark,
long[] binaryVector) {
this.scale = scale;
this.offset = offset;
@@ -456,7 +456,7 @@ public static QuantizedVector createEmpty(int quantizedDim) {
return new QuantizedVector(
Float.NaN,
Float.NaN,
- (short) 0,
+ (byte) 0,
new long[wordsForDims(quantizedDim)]
);
}
@@ -511,7 +511,7 @@ static void quantizeTo(VectorFloat> vector,
public void write(DataOutput out, int quantizedDim) throws IOException {
out.writeFloat(scale);
out.writeFloat(offset);
- out.writeInt(landmark);
+ out.writeByte(landmark);
int words = wordsForDims(quantizedDim);
for (int i = 0; i < words; i++) {
@@ -523,7 +523,12 @@ public static QuantizedVector load(RandomAccessReader in,
int quantizedDim) throws IOException {
float scale = in.readFloat();
float offset = in.readFloat();
- short landmark = (short) in.readInt();
+ byte landmark;
+ if (in instanceof io.github.jbellis.jvector.disk.ByteReadable) {
+ landmark = ((io.github.jbellis.jvector.disk.ByteReadable) in).readByte();
+ } else {
+ throw new IOException("ASH requires ByteReadable reader for landmark");
+ }
int words = wordsForDims(quantizedDim);
long[] body = new long[words];
@@ -539,7 +544,11 @@ public static void loadInto(RandomAccessReader in,
int quantizedDim) throws IOException {
dest.scale = in.readFloat();
dest.offset = in.readFloat();
- dest.landmark = (short) in.readInt();
+ if (in instanceof io.github.jbellis.jvector.disk.ByteReadable) {
+ dest.landmark = ((io.github.jbellis.jvector.disk.ByteReadable) in).readByte();
+ } else {
+ throw new IOException("ASH requires ByteReadable reader for landmark");
+ }
int words = wordsForDims(quantizedDim);
for (int i = 0; i < words; i++) {
@@ -860,7 +869,7 @@ public void encodeTo(VectorFloat> vector, QuantizedVector dest) {
// The landmark id is always 0 in this mode.
// Landmark assignment (C=1 --> always 0)
- final short landmark = 0; // placeholder; will generalize
+ final byte landmark = 0; // placeholder; will generalize
final VectorFloat> mu = landmarks[landmark];
@@ -914,7 +923,7 @@ public int compressedVectorSize() {
int words = QuantizedVector.wordsForDims(quantizedDim);
return Float.BYTES // scale
+ Float.BYTES // offset
- + Integer.BYTES // landmark
+ + Byte.BYTES // landmark
+ words * Long.BYTES; // binary payload
}
From 0d087881074db380e1c9cffc224526d532981283 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Sat, 3 Jan 2026 23:38:02 +0000
Subject: [PATCH 017/243] ASH is sensitive to SIMD alignment, so now we check
that the payload is a multiple of 64 bits.
---
.../quantization/AsymmetricHashing.java | 57 ++++++++++++++++---
.../jbellis/jvector/example/DistancesASH.java | 35 ++++++++----
.../jbellis/jvector/example/DistancesPQ.java | 23 ++++++--
3 files changed, 90 insertions(+), 25 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 5514038f7..b8762c3ea 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -175,16 +175,10 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
int optimizer,
int encodedBits) {
+ // payload bits are always 64-bit aligned by validateEncodedBits()
+ validateEncodedBits(encodedBits, HEADER_BITS);
final int quantizedDim = encodedBits - HEADER_BITS;
- if (quantizedDim < 1) {
- throw new IllegalArgumentException(
- "Illegal ASH quantized dimensionality: " + quantizedDim +
- " (encodedBits=" + encodedBits +
- ", headerBits=" + HEADER_BITS + ")"
- );
- }
-
var ravvCopy = ravv.threadLocalSupplier().get();
int originalDim = ravvCopy.getVector(0).length();
@@ -218,6 +212,52 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
);
}
+ /**
+ * Validates that the requested {@code encodedBits} value is compatible with
+ * ASH’s block-SIMD scoring kernel.
+ *
+ *
+ * For performance reasons, the binary payload size
+ * {@code (encodedBits - headerBits)} must be a multiple of 64. Non-aligned
+ * payloads introduce partial-word tails that significantly degrade SIMD
+ * efficiency.
+ *
+ *
+ * @throws IllegalArgumentException if the payload bit count is invalid or
+ * not 64-bit aligned
+ */
+ private static void validateEncodedBits(
+ int encodedBits,
+ int headerBits,
+ int originalDim
+ ) {
+ int payloadBits = encodedBits - headerBits;
+
+ if (payloadBits < 1) {
+ throw new IllegalArgumentException(
+ "Illegal ASH payload bits: " + payloadBits +
+ " (encodedBits=" + encodedBits + ", headerBits=" + headerBits + ")"
+ );
+ }
+
+ if ((payloadBits & 63) != 0) {
+ throw new IllegalArgumentException(
+ "Invalid encodedBits=" + encodedBits +
+ ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
+ "for aligned binary payload. " +
+ "Got payloadBits=" + payloadBits +
+ " (headerBits=" + headerBits + ")."
+ );
+ }
+
+ if (payloadBits > originalDim) {
+ throw new IllegalArgumentException(
+ "Invalid ASH payloadBits=" + payloadBits +
+ " exceeds originalDim=" + originalDim
+ );
+ }
+ }
+
/**
* Loads an ASH instance from a RandomAccessReader.
* This must mirror the ASH header writer exactly.
@@ -238,6 +278,7 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
int originalDimension = in.readInt();
int encodedBits = in.readInt();
+ validateEncodedBits(encodedBits, HEADER_BITS, originalDimension);
int quantizedDim = in.readInt();
int optimizer = in.readInt();
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index d61d904c5..801f204c5 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -52,7 +52,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "true"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
final boolean RUN_SCALAR_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
@@ -60,12 +60,15 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
final boolean RUN_FLOAT_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
+ // ASH header bits
+ final int HEADER_BITS = 72;
+
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {16, 32, 64}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {64}; // 16, 32, and/or 64
// Define the benchmark size
- int maxQueries = 1_000;
- int maxVectors = 100_000;
+ int maxQueries = 10_000;
+ int maxVectors = 10_000_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -119,7 +122,17 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = dimension / 4; // sweep later if desired
+ int encodedBits = (dimension / 4) + HEADER_BITS;
+ // Payload must be 64-bit aligned for SIMD
+ int payloadBits = encodedBits - HEADER_BITS;
+ if ((payloadBits & 63) != 0) {
+ throw new IllegalArgumentException(
+ "ASH payloadBits must be 64-bit aligned for SIMD. " +
+ "Got payloadBits=" + payloadBits +
+ " (encodedBits=" + encodedBits +
+ ", HEADER_BITS=" + HEADER_BITS + ")"
+ );
+ }
System.out.println(
"\toriginalDim = " + dimension +
@@ -493,11 +506,11 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
- runCohere100k();
- runADA();
- runOpenai1536();
- runOpenai3072();
-// runCap6m();
-// runCohere10m();
+// runCohere100k();
+// runADA();
+// runOpenai1536();
+// runOpenai3072();
+ runCap6m();
+ runCohere10m();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
index 6c5fe9190..47c88f975 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -121,11 +121,21 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
int chunkSize = Math.max(1, (finalQueryCount + parallelism - 1) / parallelism);
// ------------------------------------------------------------
- // PQ parameters (compression matched to ASH)
+ // PQ parameters (EXACT bit-for-bit match with ASH)
// ------------------------------------------------------------
- // ASH: encodedBits = d / 4 -> bytes = d / 32
- // PQ: 1 byte per subspace -> M = d / 32
- final int M = Math.max(1, dimension / 32);
+ // ASH total size = encodedBits (includes header + payload)
+ // PQ total size = M * 8 bits (1 byte per subspace)
+ //
+ // We REQUIRE exact equality for fair benchmarking.
+ int HEADER_BITS = 72;
+ int encodedBitsASH = (dimension / 4) + HEADER_BITS; // HEADER_BITS = 72
+ if ((encodedBitsASH & 7) != 0) {
+ throw new IllegalArgumentException(
+ "PQ requires encodedBits to be a multiple of 8 for exact bit matching; got " + encodedBitsASH
+ );
+ }
+
+ final int M = encodedBitsASH / 8;
final int K = 256; // 1 byte per subspace
final boolean globallyCenter = false;
final float anisotropicThreshold = 0.0f;
@@ -133,8 +143,9 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
System.out.println("\tPQ params:");
System.out.println("\t M = " + M + " subspaces");
System.out.println("\t K = " + K + " clusters");
- System.out.println("\t Compression ≈ " + (M * 8) + " bits (" + M + " bytes)");
- System.out.println("\t ASH reference ≈ " + (dimension / 4) + " bits (" + (dimension / 32) + " bytes)");
+ System.out.println("\t Compression = " + (M * 8) + " bits (" + M + " bytes)");
+ System.out.println("\tASH reference = " + encodedBitsASH + " bits (" + (encodedBitsASH / 8) + " bytes)");
+
// ------------------------------------------------------------
// Train PQ + encode (timed)
From 19a3534b902e4ca33077c99817a1ed18b18a740f Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Sat, 3 Jan 2026 23:38:02 +0000
Subject: [PATCH 018/243] ASH is sensitive to SIMD alignment, so now we check
that the payload is a multiple of 64 bits.
---
.../quantization/AsymmetricHashing.java | 57 ++++++++++++++++---
.../jbellis/jvector/example/DistancesASH.java | 35 ++++++++----
.../jbellis/jvector/example/DistancesPQ.java | 23 ++++++--
3 files changed, 90 insertions(+), 25 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 5514038f7..bb45c9158 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -177,16 +177,10 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
final int quantizedDim = encodedBits - HEADER_BITS;
- if (quantizedDim < 1) {
- throw new IllegalArgumentException(
- "Illegal ASH quantized dimensionality: " + quantizedDim +
- " (encodedBits=" + encodedBits +
- ", headerBits=" + HEADER_BITS + ")"
- );
- }
-
var ravvCopy = ravv.threadLocalSupplier().get();
int originalDim = ravvCopy.getVector(0).length();
+ // payload bits are always 64-bit aligned by validateEncodedBits()
+ validateEncodedBits(encodedBits, HEADER_BITS, originalDim);
// C=1 landmarks case
VectorFloat> mu0 = vectorTypeSupport.createFloatVector(originalDim);
@@ -218,6 +212,52 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
);
}
+ /**
+ * Validates that the requested {@code encodedBits} value is compatible with
+ * ASH’s block-SIMD scoring kernel.
+ *
+ *
+ * For performance reasons, the binary payload size
+ * {@code (encodedBits - headerBits)} must be a multiple of 64. Non-aligned
+ * payloads introduce partial-word tails that significantly degrade SIMD
+ * efficiency.
+ *
+ *
+ * @throws IllegalArgumentException if the payload bit count is invalid or
+ * not 64-bit aligned
+ */
+ private static void validateEncodedBits(
+ int encodedBits,
+ int headerBits,
+ int originalDim
+ ) {
+ int payloadBits = encodedBits - headerBits;
+
+ if (payloadBits < 1) {
+ throw new IllegalArgumentException(
+ "Illegal ASH payload bits: " + payloadBits +
+ " (encodedBits=" + encodedBits + ", headerBits=" + headerBits + ")"
+ );
+ }
+
+ if ((payloadBits & 63) != 0) {
+ throw new IllegalArgumentException(
+ "Invalid encodedBits=" + encodedBits +
+ ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
+ "for aligned binary payload. " +
+ "Got payloadBits=" + payloadBits +
+ " (headerBits=" + headerBits + ")."
+ );
+ }
+
+ if (payloadBits > originalDim) {
+ throw new IllegalArgumentException(
+ "Invalid ASH payloadBits=" + payloadBits +
+ " exceeds originalDim=" + originalDim
+ );
+ }
+ }
+
/**
* Loads an ASH instance from a RandomAccessReader.
* This must mirror the ASH header writer exactly.
@@ -238,6 +278,7 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
int originalDimension = in.readInt();
int encodedBits = in.readInt();
+ validateEncodedBits(encodedBits, HEADER_BITS, originalDimension);
int quantizedDim = in.readInt();
int optimizer = in.readInt();
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index d61d904c5..801f204c5 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -52,7 +52,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "true"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
final boolean RUN_SCALAR_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
@@ -60,12 +60,15 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
final boolean RUN_FLOAT_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
+ // ASH header bits
+ final int HEADER_BITS = 72;
+
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {16, 32, 64}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {64}; // 16, 32, and/or 64
// Define the benchmark size
- int maxQueries = 1_000;
- int maxVectors = 100_000;
+ int maxQueries = 10_000;
+ int maxVectors = 10_000_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -119,7 +122,17 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = dimension / 4; // sweep later if desired
+ int encodedBits = (dimension / 4) + HEADER_BITS;
+ // Payload must be 64-bit aligned for SIMD
+ int payloadBits = encodedBits - HEADER_BITS;
+ if ((payloadBits & 63) != 0) {
+ throw new IllegalArgumentException(
+ "ASH payloadBits must be 64-bit aligned for SIMD. " +
+ "Got payloadBits=" + payloadBits +
+ " (encodedBits=" + encodedBits +
+ ", HEADER_BITS=" + HEADER_BITS + ")"
+ );
+ }
System.out.println(
"\toriginalDim = " + dimension +
@@ -493,11 +506,11 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
- runCohere100k();
- runADA();
- runOpenai1536();
- runOpenai3072();
-// runCap6m();
-// runCohere10m();
+// runCohere100k();
+// runADA();
+// runOpenai1536();
+// runOpenai3072();
+ runCap6m();
+ runCohere10m();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
index 6c5fe9190..47c88f975 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -121,11 +121,21 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
int chunkSize = Math.max(1, (finalQueryCount + parallelism - 1) / parallelism);
// ------------------------------------------------------------
- // PQ parameters (compression matched to ASH)
+ // PQ parameters (EXACT bit-for-bit match with ASH)
// ------------------------------------------------------------
- // ASH: encodedBits = d / 4 -> bytes = d / 32
- // PQ: 1 byte per subspace -> M = d / 32
- final int M = Math.max(1, dimension / 32);
+ // ASH total size = encodedBits (includes header + payload)
+ // PQ total size = M * 8 bits (1 byte per subspace)
+ //
+ // We REQUIRE exact equality for fair benchmarking.
+ int HEADER_BITS = 72;
+ int encodedBitsASH = (dimension / 4) + HEADER_BITS; // HEADER_BITS = 72
+ if ((encodedBitsASH & 7) != 0) {
+ throw new IllegalArgumentException(
+ "PQ requires encodedBits to be a multiple of 8 for exact bit matching; got " + encodedBitsASH
+ );
+ }
+
+ final int M = encodedBitsASH / 8;
final int K = 256; // 1 byte per subspace
final boolean globallyCenter = false;
final float anisotropicThreshold = 0.0f;
@@ -133,8 +143,9 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
System.out.println("\tPQ params:");
System.out.println("\t M = " + M + " subspaces");
System.out.println("\t K = " + K + " clusters");
- System.out.println("\t Compression ≈ " + (M * 8) + " bits (" + M + " bytes)");
- System.out.println("\t ASH reference ≈ " + (dimension / 4) + " bits (" + (dimension / 32) + " bytes)");
+ System.out.println("\t Compression = " + (M * 8) + " bits (" + M + " bytes)");
+ System.out.println("\tASH reference = " + encodedBitsASH + " bits (" + (encodedBitsASH / 8) + " bytes)");
+
// ------------------------------------------------------------
// Train PQ + encode (timed)
From ae8bf0434147d98cd6ebd4e413ab5e585d2c92f0 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Sun, 4 Jan 2026 19:24:28 +0000
Subject: [PATCH 019/243] Added multi-landmark support to the scalar scorer.
Updated DistancesASH to acceptt landmark count.
---
.../jvector/quantization/ASHScorer.java | 123 +++++----
.../quantization/AsymmetricHashing.java | 256 ++++++++++++++++--
.../jbellis/jvector/example/DistancesASH.java | 7 +-
3 files changed, 313 insertions(+), 73 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
index 7ba87b284..a345a088f 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHScorer.java
@@ -22,7 +22,11 @@
/**
* Query-time scorer for ASH vectors.
- * Single-landmark case (C = 1).
+ *
+ * Multi-landmark scoring (C >= 1) following the paper’s Eq. 11 decomposition:
+ * for each landmark c we form q̃_c = A (q - μ_c) = (Aq) - (Aμ_c) once per query,
+ * then score each vector using its stored landmark id.
+ *
* NOTE: Only DOT_PRODUCT is supported for now.
*/
public final class ASHScorer {
@@ -30,6 +34,22 @@ public final class ASHScorer {
public ASHScorer(AsymmetricHashing ash) {
this.ash = ash;
+
+ final int d = ash.quantizedDim;
+ final int C = ash.landmarkCount;
+
+ assert ash.landmarkProj != null;
+ assert ash.landmarkProjSum != null;
+ assert ash.landmarkProj.length == C :
+ "landmarkProj length mismatch";
+ assert ash.landmarkProjSum.length == C :
+ "landmarkProjSum length mismatch";
+
+ for (int c = 0; c < C; c++) {
+ assert ash.landmarkProj[c].length == d :
+ "landmarkProj[" + c + "] length != quantizedDim";
+ }
+
assert ash.stiefelTransform.rows == ash.quantizedDim
: "Stiefel rows must equal quantizedDim";
}
@@ -48,80 +68,89 @@ public ASHScoreFunction scoreFunctionFor(VectorFloat> query, VectorSimilarityF
*/
private ASHScoreFunction dotProductScoreFunctionFor(VectorFloat> query) {
final int d = ash.quantizedDim;
+ final int D = ash.originalDimension;
+ final int C = ash.landmarkCount;
- // Compute q_centered = q - μ
- // (avoid allocation by doing subtraction on the fly during projection)
- // Landmark-dependent query projection (C=1 => still index 0)
- final VectorFloat> mu = ash.landmarks[0];
-
- final float[][] A = ash.stiefelTransform.AFloat; // shape: d × originalDim
- final int originalDim = ash.stiefelTransform.cols;
- final float[] muArr = new float[originalDim];
+ final float[][] A = ash.stiefelTransform.AFloat; // [d][D]
final var vecUtil =
io.github.jbellis.jvector.vector.VectorizationProvider
.getInstance()
.getVectorUtilSupport();
- for (int k = 0; k < originalDim; k++) muArr[k] = mu.get(k);
-
- final float[] qArr = new float[originalDim];
- for (int k = 0; k < originalDim; k++) qArr[k] = query.get(k);
+ // Materialize query once (avoid VectorFloat.get in inner loops)
+ final float[] qArr = new float[D];
+ for (int k = 0; k < D; k++) qArr[k] = query.get(k);
- // Project query: tildeQ[j] = (A (q - μ))[j]
- final float[] tildeQ = new float[d];
+ // Compute Aq once per query:
+ // qProj[j] =
+ final float[] qProj = new float[d];
+ float sumQProj = 0f;
for (int j = 0; j < d; j++) {
- final float[] Aj = A[j];
- float acc = 0.0f;
- for (int k = 0; k < originalDim; k++) {
- acc += Aj[k] * (qArr[k] - muArr[k]);
- }
- tildeQ[j] = acc;
+ float v = vecUtil.ashDotRow(A[j], qArr);
+ qProj[j] = v;
+ sumQProj += v;
}
- // = sum_j tildeQ[j]
- float tmpSum = 0f;
- for (int j = 0; j < d; j++) {
- tmpSum += tildeQ[j];
+ // For each landmark c:
+ // tildeQ_c = Aq - Aμ_c
+ // sumTildeQc = sum(Aq) - sum(Aμ_c)
+ // dotQMu_c =
+ // Pooled storage for q̃_c vectors:
+ // tildeQPool[c*d + j] == q̃_c[j]
+ final float[] tildeQPool = new float[C * d];
+ final float[] sumTildeQByLandmark = new float[C];
+ final float[] dotQMuByLandmark = new float[C];
+
+ for (int c = 0; c < C; c++) {
+ dotQMuByLandmark[c] = VectorUtil.dotProduct(query, ash.landmarks[c]);
+ sumTildeQByLandmark[c] = sumQProj - ash.landmarkProjSum[c];
+
+ final int base = c * d;
+ final float[] muProj = ash.landmarkProj[c]; // [d]
+ for (int j = 0; j < d; j++) {
+ tildeQPool[base + j] = qProj[j] - muProj[j];
+ }
}
- final float sumTildeQ = tmpSum;
-
- //
- final float dotQMu = VectorUtil.dotProduct(query, mu);
return (AsymmetricHashing.QuantizedVector v) -> {
- // Single landmark baseline: ignore v.landmark for now (assumed 0)
- final float scale = v.scale; // already ||x − μ|| / sqrt(d)
- final float offset = v.offset; // already − ||μ||^2
+ final int c = v.landmark & 0xFF; // unsigned [0,C)
+ assert c < C : "Invalid landmark id " + c + " for landmarkCount=" + C;
- // masked-add term: where bin ∈ {0,1}^d
- // This is sum of tildeQ[j] over set bits in v.binaryVector.
+ // scale = ||x − μ_c|| / √d (precomputed at encode-time)
+ final float scale = v.scale;
+
+ // offset = − ||μ_c||² (precomputed at encode-time)
+ final float offset = v.offset;
+
+ // maskedAdd =
+ // b ∈ {0,1}^d stored as packed longs
float maskedAdd = 0f;
- // v.binaryVector stores bits packed in longs, little-endian per word (bit j corresponds to dimension base+j)
+ final int tildeBase = c * d;
+ final float sumTildeQ = sumTildeQByLandmark[c];
+ final float dotQMu = dotQMuByLandmark[c];
+
final long[] bits = v.binaryVector;
- int base = 0;
- for (int w = 0; w < bits.length && base < d; w++, base += 64) {
+ int dimBase = 0;
+ for (int w = 0; w < bits.length && dimBase < d; w++, dimBase += 64) {
long word = bits[w];
while (word != 0L) {
int bit = Long.numberOfTrailingZeros(word);
- int idx = base + bit;
- if (idx < d) maskedAdd += tildeQ[idx];
+ int idx = dimBase + bit;
+ if (idx < d) {
+ maskedAdd += tildeQPool[tildeBase + idx];
+ }
word &= (word - 1);
}
}
- // ASH dot-product approximation (single landmark, {0,1} encoding):
+ // ASH dot-product approximation (multi-landmark, {0,1} encoding):
//
- // scale * (2⟨q̃, b⟩ − ⟨q̃, 1⟩) + ⟨q, μ⟩ + (⟨x, μ⟩ − ||μ||²)
+ // scale * (2⟨q̃_c, b⟩ − ⟨q̃_c, 1⟩) + ⟨q, μ_c⟩ + (⟨x, μ_c⟩ − ||μ_c||²)
//
- // Notes:
- // - b ∈ {0,1}^d is the binary code (converted from the paper’s {−1,+1} form).
- // - scale = ||x − μ|| / √d (pure Eq. 11 baseline; no bias correction).
- // - ⟨q, μ⟩ is a query-only offset and MUST NOT be scaled.
- // - Scaling ⟨q, μ⟩ leads to large systematic error (see paper Eq. 11 regrouping).
+ // where q̃_c = A(q − μ_c) is precomputed once per query and per landmark c.
return scale * (2f * maskedAdd - sumTildeQ) + dotQMu + offset;
-
};
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index bb45c9158..5f253d4f8 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -90,6 +90,22 @@ public class AsymmetricHashing implements VectorCompressor where A is [d][D] and μ_c is length D.
+ *
+ * Shape: [landmarkCount][quantizedDim]
+ */
+ public final float[][] landmarkProj;
+
+ /**
+ * Precomputed sums of the landmark projections:
+ * landmarkProjSum[c] = Σ_j landmarkProj[c][j]
+ *
+ * Shape: [landmarkCount]
+ */
+ public final float[] landmarkProjSum;
+
/** Debug hook to disable learning paths */
@VisibleForTesting
public boolean learn = true;
@@ -125,10 +141,22 @@ private AsymmetricHashing(int originalDim,
int quantizedDim,
int optimizer,
VectorFloat>[] landmarks,
- StiefelTransform stiefelTransform) {
+ StiefelTransform stiefelTransform,
+ float[][] landmarkProj,
+ float[] landmarkProjSum) {
this.originalDimension = originalDim;
this.encodedBits = encodedBits;
this.quantizedDim = quantizedDim;
+
+ // Defensive check for consistency
+ if (encodedBits - HEADER_BITS != quantizedDim) {
+ throw new IllegalArgumentException(
+ "Invalid ASH configuration: encodedBits=" + encodedBits +
+ ", quantizedDim=" + quantizedDim +
+ ", HEADER_BITS=" + HEADER_BITS
+ );
+ }
+
this.optimizer = optimizer;
this.stiefelTransform = Objects.requireNonNull(stiefelTransform);
@@ -143,6 +171,25 @@ private AsymmetricHashing(int originalDim,
VectorUtil.dotProduct(landmarks[c], landmarks[c]);
}
+ this.landmarkProj = Objects.requireNonNull(landmarkProj, "landmarkProj");
+ this.landmarkProjSum = Objects.requireNonNull(landmarkProjSum, "landmarkProjSum");
+
+ if (landmarkProj.length != landmarkCount) {
+ throw new IllegalArgumentException("landmarkProj outer dimension mismatch: "
+ + landmarkProj.length + " != landmarkCount=" + landmarkCount);
+ }
+ if (landmarkProjSum.length != landmarkCount) {
+ throw new IllegalArgumentException("landmarkProjSum length mismatch: "
+ + landmarkProjSum.length + " != landmarkCount=" + landmarkCount);
+ }
+ for (int c = 0; c < landmarkCount; c++) {
+ if (landmarkProj[c] == null || landmarkProj[c].length != quantizedDim) {
+ throw new IllegalArgumentException("landmarkProj[" + c + "] length mismatch: "
+ + (landmarkProj[c] == null ? "null" : landmarkProj[c].length)
+ + " != quantizedDim=" + quantizedDim);
+ }
+ }
+
if (landmarkNormSq.length != landmarkCount) {
throw new IllegalArgumentException("landmarkNormSq length mismatch");
}
@@ -173,23 +220,65 @@ private AsymmetricHashing(int originalDim,
*/
public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
int optimizer,
- int encodedBits) {
+ int encodedBits,
+ int landmarkCount) {
final int quantizedDim = encodedBits - HEADER_BITS;
+ if (landmarkCount < 1 || landmarkCount > 64) {
+ throw new IllegalArgumentException(
+ "landmarkCount must be in [1,64], got " + landmarkCount);
+ }
+
+ System.out.println(
+ "\tASH initialized with landmarkCount=" + landmarkCount +
+ ", quantizedDim=" + quantizedDim
+ );
+
var ravvCopy = ravv.threadLocalSupplier().get();
int originalDim = ravvCopy.getVector(0).length();
// payload bits are always 64-bit aligned by validateEncodedBits()
validateEncodedBits(encodedBits, HEADER_BITS, originalDim);
- // C=1 landmarks case
- VectorFloat> mu0 = vectorTypeSupport.createFloatVector(originalDim);
+ // NOTE: points are treated as read-only by KMeansPlusPlusClusterer. Materialize once.
+ VectorFloat>[] points = new VectorFloat>[ravvCopy.size()];
for (int i = 0; i < ravvCopy.size(); i++) {
- VectorUtil.addInPlace(mu0, ravvCopy.getVector(i));
+ points[i] = ravvCopy.getVector(i);
}
- VectorUtil.scale(mu0, 1.0f / ravvCopy.size());
- VectorFloat>[] landmarks = new VectorFloat>[] { mu0 };
+ VectorFloat>[] landmarks;
+ if (landmarkCount == 1) {
+ // Mean centroid
+ VectorFloat> mu0 = vectorTypeSupport.createFloatVector(originalDim);
+ for (int i = 0; i < points.length; i++) {
+ VectorUtil.addInPlace(mu0, points[i]);
+ }
+ VectorUtil.scale(mu0, 1.0f / points.length);
+ landmarks = new VectorFloat>[] { mu0 };
+ } else {
+ long kmStart = System.nanoTime();
+
+ // KMeans++ in JVector returns centroids packed as one big vector [C * D]
+ KMeansPlusPlusClusterer km = new KMeansPlusPlusClusterer(points, landmarkCount);
+ VectorFloat> packed = km.cluster(/*unweightedIterations*/ 20, /*anisotropicIterations*/ 0);
+
+ long kmEnd = System.nanoTime();
+
+ System.out.printf(
+ "\tKMeans++ (C=%d, N=%d, D=%d) took %.3f seconds%n",
+ landmarkCount,
+ points.length,
+ originalDim,
+ (kmEnd - kmStart) / 1e9
+ );
+
+ landmarks = new VectorFloat>[landmarkCount];
+ for (int c = 0; c < landmarkCount; c++) {
+ VectorFloat> mu = vectorTypeSupport.createFloatVector(originalDim);
+ mu.copyFrom(packed, c * originalDim, 0, originalDim);
+ landmarks[c] = mu;
+ }
+ }
StiefelTransform stiefelTransform;
if (optimizer == RANDOM) {
@@ -202,13 +291,23 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
throw new IllegalArgumentException("Unknown optimizer " + optimizer);
}
+ LandmarkProjections lp =
+ computeLandmarkProjections(
+ landmarks,
+ stiefelTransform,
+ quantizedDim,
+ originalDim
+ );
+
return new AsymmetricHashing(
originalDim,
encodedBits,
quantizedDim,
optimizer,
landmarks,
- stiefelTransform
+ stiefelTransform,
+ lp.proj,
+ lp.sum
);
}
@@ -258,6 +357,67 @@ private static void validateEncodedBits(
}
}
+ /**
+ * Precompute landmark projections in the reduced space:
+ *
+ * landmarkProj[c][j] =
+ * landmarkProjSum[c] = Σ_j landmarkProj[c][j]
+ *
+ * This is used by both initialize() and load() to ensure
+ * identical query-time math regardless of construction path.
+ */
+ private static LandmarkProjections computeLandmarkProjections(
+ VectorFloat>[] landmarks,
+ StiefelTransform stiefelTransform,
+ int quantizedDim,
+ int originalDim
+ ) {
+ final int C = landmarks.length;
+ final float[][] A = stiefelTransform.AFloat; // [d][D]
+
+ final var vecUtil =
+ io.github.jbellis.jvector.vector.VectorizationProvider
+ .getInstance()
+ .getVectorUtilSupport();
+
+ final float[][] landmarkProj = new float[C][quantizedDim];
+ final float[] landmarkProjSum = new float[C];
+
+ final float[] muArr = new float[originalDim];
+
+ for (int c = 0; c < C; c++) {
+ VectorFloat> mu = landmarks[c];
+
+ // materialize μ once
+ for (int i = 0; i < originalDim; i++) {
+ muArr[i] = mu.get(i);
+ }
+
+ float sum = 0f;
+ float[] proj = landmarkProj[c];
+
+ for (int j = 0; j < quantizedDim; j++) {
+ float v = vecUtil.ashDotRow(A[j], muArr);
+ proj[j] = v;
+ sum += v;
+ }
+
+ landmarkProjSum[c] = sum;
+ }
+
+ return new LandmarkProjections(landmarkProj, landmarkProjSum);
+ }
+
+ private static final class LandmarkProjections {
+ final float[][] proj; // [C][d]
+ final float[] sum; // [C]
+
+ LandmarkProjections(float[][] proj, float[] sum) {
+ this.proj = proj;
+ this.sum = sum;
+ }
+ }
+
/**
* Loads an ASH instance from a RandomAccessReader.
* This must mirror the ASH header writer exactly.
@@ -280,14 +440,36 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
int encodedBits = in.readInt();
validateEncodedBits(encodedBits, HEADER_BITS, originalDimension);
int quantizedDim = in.readInt();
+
+ // HEADER_BITS compatibility check
+ int expectedQuantizedDim = encodedBits - HEADER_BITS;
+ if (quantizedDim != expectedQuantizedDim) {
+ throw new IOException(
+ "ASH format mismatch: encodedBits=" + encodedBits +
+ ", quantizedDim=" + quantizedDim +
+ ", but runtime HEADER_BITS=" + HEADER_BITS +
+ " implies quantizedDim=" + expectedQuantizedDim
+ );
+ }
+
int optimizer = in.readInt();
- int landmark0Length = in.readInt();
- if (landmark0Length <= 0) {
- throw new IOException("ASH header missing landmark[0]");
+ // Multi-landmark support
+ int landmarkCount = in.readInt();
+ if (landmarkCount < 1 || landmarkCount > 64) {
+ throw new IOException("Invalid landmarkCount=" + landmarkCount);
+ }
+
+ VectorFloat>[] landmarks = new VectorFloat>[landmarkCount];
+ for (int c = 0; c < landmarkCount; c++) {
+ int dim = in.readInt();
+ if (dim != originalDimension) {
+ throw new IOException(
+ "Landmark dimension mismatch: expected " +
+ originalDimension + ", got " + dim);
+ }
+ landmarks[c] = vectorTypeSupport.readFloatVector(in, dim);
}
- VectorFloat> mu0 = vectorTypeSupport.readFloatVector(in, landmark0Length);
- VectorFloat>[] landmarks = new VectorFloat>[] { mu0 };
// NOTE: StiefelTransform loading will be added in Step 2.
// For now, construct a placeholder or re-run the RANDOM initializer.
@@ -306,13 +488,23 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
);
}
+ LandmarkProjections lp =
+ computeLandmarkProjections(
+ landmarks,
+ stiefelTransform,
+ quantizedDim,
+ originalDimension
+ );
+
return new AsymmetricHashing(
originalDimension,
encodedBits,
quantizedDim,
optimizer,
landmarks,
- stiefelTransform
+ stiefelTransform,
+ lp.proj,
+ lp.sum
);
}
@@ -337,9 +529,13 @@ public void write(DataOutput out, int version) throws IOException {
out.writeInt(quantizedDim);
out.writeInt(optimizer);
- // (C=1): write landmark[0] (dataset mean). Multi-landmark serialization TBD.
- out.writeInt(landmarks[0].length());
- vectorTypeSupport.writeFloatVector(out, landmarks[0]);
+ // Multi-landmark serialization
+ out.writeInt(landmarkCount);
+
+ for (int c = 0; c < landmarkCount; c++) {
+ out.writeInt(landmarks[c].length());
+ vectorTypeSupport.writeFloatVector(out, landmarks[c]);
+ }
// NOTE: stiefelTransform serialization will be added later
}
@@ -350,7 +546,7 @@ public void write(DataOutput out, int version) throws IOException {
public void quantizeVector(VectorFloat> vector,
float residualNorm,
- short landmark,
+ byte landmark,
QuantizedVector dest) {
if (quantizedDim > originalDimension) {
@@ -503,7 +699,7 @@ public static QuantizedVector createEmpty(int quantizedDim) {
}
public static int serializedSizeBytes(int quantizedDim) {
- return Float.BYTES + Float.BYTES + Short.BYTES
+ return Float.BYTES + Float.BYTES + Byte.BYTES
+ wordsForDims(quantizedDim) * Long.BYTES;
}
@@ -909,10 +1105,19 @@ public void encodeTo(VectorFloat> vector, QuantizedVector dest) {
// All vectors are encoded relative to landmark μ_c (C=1 => μ_0 is the dataset mean).
// The landmark id is always 0 in this mode.
- // Landmark assignment (C=1 --> always 0)
- final byte landmark = 0; // placeholder; will generalize
+ // Landmark assignment: nearest centroid by L2 distance
+ byte landmark = 0;
+ float bestDist = Float.POSITIVE_INFINITY;
- final VectorFloat> mu = landmarks[landmark];
+ for (int c = 0; c < landmarkCount; c++) {
+ float dist = VectorUtil.squareL2Distance(vector, landmarks[c]);
+ if (dist < bestDist) {
+ bestDist = dist;
+ landmark = (byte) c;
+ }
+ }
+
+ final VectorFloat> mu = landmarks[landmark & 0xFF];
// Compute
final float dotXMu = VectorUtil.dotProduct(vector, mu);
@@ -954,8 +1159,11 @@ public int compressorSize() {
size += Integer.BYTES; // encodedBits
size += Integer.BYTES; // quantizedDim
size += Integer.BYTES; // optimizer
- size += Integer.BYTES; // landmarks[0] length
- size += Float.BYTES * landmarks[0].length();
+ size += Integer.BYTES; // landmarkCount
+ for (int c = 0; c < landmarkCount; c++) {
+ size += Integer.BYTES; // landmark dimension
+ size += Float.BYTES * landmarks[c].length();
+ }
return size;
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 801f204c5..37f0a20cf 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -52,7 +52,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "true"));
final boolean RUN_SCALAR_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
@@ -66,6 +66,9 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Block sizes to benchmark
final int[] BLOCK_SIZES = {64}; // 16, 32, and/or 64
+ // How many ASH landmarks to use, C = [1, 64]
+ final int landmarkCount = 64;
+
// Define the benchmark size
int maxQueries = 10_000;
int maxVectors = 10_000_000;
@@ -149,7 +152,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Build ASH
// ------------------------------------------------------------------
var ravv = new ListRandomAccessVectorValues(vectors, dimension);
- var ash = AsymmetricHashing.initialize(ravv, AsymmetricHashing.RANDOM, encodedBits);
+ var ash = AsymmetricHashing.initialize(ravv, AsymmetricHashing.RANDOM, encodedBits, landmarkCount);
long startTime = System.nanoTime();
CompressedVectors ashVecs = ash.encodeAll(ravv);
From 92bf98d4ed431bc1dd8b2a10c3ec2eacc9db9f35 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Mon, 5 Jan 2026 18:05:45 +0000
Subject: [PATCH 020/243] Created a pooled landmark-specific query buffer.
Still using the non-pooled version for performance purposes though.
---
.../jvector/quantization/ASHVectors.java | 265 ++++++++++--------
.../jvector/vector/VectorUtilSupport.java | 51 ++++
.../vector/PanamaVectorUtilSupport.java | 67 +++++
3 files changed, 272 insertions(+), 111 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index 4e6bc43df..57ea8debf 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -169,9 +169,9 @@ public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(
return node -> f.similarityTo(compressedVectors[node]);
}
-// ============================================================================
-// ASH block-oriented scoring
-// ============================================================================
+ // ============================================================================
+ // ASH block-oriented scoring
+ // ============================================================================
/**
* Scalar reference block scorer.
@@ -282,16 +282,19 @@ public ASHBlockScorer blockScorerFor(
final AshBlockKernel kernel = AshBlockKernel.fromProperty();
+ // Compute query precompute ONCE (shared by scalar + simd)
+ final QueryPrecompute qp = precomputeQuery(query);
+
// ------------------------------------------------------------
// Forced scalar path (portable register-accumulator)
// ------------------------------------------------------------
if (kernel == AshBlockKernel.SCALAR) {
return new BlockScalarASHBlockScorer(
- ash,
compressedVectors,
scales,
offsets,
- query
+ landmarks,
+ qp
);
}
@@ -304,13 +307,13 @@ public ASHBlockScorer blockScorerFor(
ensurePackedBits(blockSize);
return new BlockSimdASHBlockScorer(
- ash,
compressedVectors,
scales,
offsets,
+ landmarks,
packedBits,
blockSize,
- query
+ qp
);
}
@@ -354,89 +357,72 @@ private void ensurePackedBits(int blockSize) {
*/
private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
- private final AsymmetricHashing ash;
private final AsymmetricHashing.QuantizedVector[] vectors;
private final float[] scales;
private final float[] offsets;
+ private final byte[] landmarks;
-
- private final float[] tildeQ;
- private final float sumTildeQ;
- private final float dotQMu;
-
+ private final float[] tildeQPool; // [C * d]
+ private final float[] sumTildeQ; // [C]
+ private final float[] dotQMu; // [C]
private final int d;
private final int words;
BlockScalarASHBlockScorer(
- AsymmetricHashing ash,
AsymmetricHashing.QuantizedVector[] vectors,
float[] scales,
float[] offsets,
- VectorFloat> query) {
+ byte[] landmarks,
+ QueryPrecompute qp)
+ {
- this.ash = ash;
this.vectors = vectors;
this.scales = scales;
this.offsets = offsets;
+ this.landmarks = landmarks;
- this.d = ash.quantizedDim;
+ this.d = qp.d;
this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
- final int D = ash.originalDimension;
- final VectorFloat> mu = ash.landmarks[0];
-
- final float[] qArr = new float[D];
- final float[] muArr = new float[D];
- for (int i = 0; i < D; i++) {
- qArr[i] = query.get(i);
- muArr[i] = mu.get(i);
- }
-
- this.tildeQ = new float[d];
- final float[][] A = ash.stiefelTransform.AFloat;
- for (int r = 0; r < d; r++) {
- float s = 0f;
- float[] Arow = A[r];
- for (int j = 0; j < D; j++) {
- s += Arow[j] * (qArr[j] - muArr[j]);
- }
- tildeQ[r] = s;
- }
-
- float sum = 0f;
- for (int i = 0; i < d; i++) sum += tildeQ[i];
- this.sumTildeQ = sum;
- this.dotQMu = VectorUtil.dotProduct(query, mu);
+ this.tildeQPool = qp.tildeQPool;
+ this.sumTildeQ = qp.sumTildeQ;
+ this.dotQMu = qp.dotQMu;
}
@Override
public void scoreRange(int start, int count, float[] out) {
+ if (count < 0) throw new IllegalArgumentException("count must be >= 0");
+ if (start < 0 || start + count > vectors.length) {
+ throw new IllegalArgumentException("Range out of bounds: [" + start + ", " + (start + count) + ")");
+ }
+ if (out.length < count) throw new IllegalArgumentException("out.length < count");
+ if (count == 0) return;
for (int i = 0; i < count; i++) {
- int idx = start + i;
+ final int idx = start + i;
+
+ final int c = landmarks[idx] & 0xFF; // unsigned [0,C)
+ final int base = c * d;
float maskedAdd = 0f;
- long[] bits = vectors[idx].binaryVector;
- for (int w = 0; w < words; w++) {
+ // bit-walk against landmark-specific q̃
+ final long[] bits = vectors[idx].binaryVector;
+ int baseDim = 0;
+ for (int w = 0; w < words && baseDim < d; w++, baseDim += 64) {
long word = bits[w];
- int baseDim = w * 64;
-
while (word != 0L) {
int bit = Long.numberOfTrailingZeros(word);
- int dim = baseDim + bit;
- if (dim < d) maskedAdd += tildeQ[dim];
+ int j = baseDim + bit;
+ if (j < d) maskedAdd += tildeQPool[base + j];
word &= (word - 1);
}
}
- float scale = scales[idx];
- float offset = offsets[idx];
-
out[i] =
- scale * (2f * maskedAdd - sumTildeQ)
- + dotQMu
- + offset;
+ scales[idx] * (2f * maskedAdd - sumTildeQ[c])
+ + dotQMu[c]
+ + offsets[idx];
}
}
}
@@ -464,11 +450,10 @@ public void scoreRange(int start, int count, float[] out) {
*/
private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
- private final AsymmetricHashing ash;
private final AsymmetricHashing.QuantizedVector[] vectors;
-
private final float[] scales;
private final float[] offsets;
+ private final byte[] landmarks;
private final long[] packedBits;
private final int blockSize;
@@ -476,115 +461,173 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private final int d;
private final int words;
- // Query precompute
- private final float[] tildeQ;
- private final float sumTildeQ;
- private final float dotQMu;
+ private final float[] maskedAdds; // [blockSize]
- // Scratch (one per scorer instance)
- private final float[] maskedAdds;
+ private final float[] tildeQPool; // qp.tildeQPool [C*d]
+ private final float[] tildeQScratch; // [d]
+ private final float[] sumTildeQ; // qp.sumTildeQ [C]
+ private final float[] dotQMu; // qp.dotQMu [C]
private final VectorUtilSupport vecUtil;
BlockSimdASHBlockScorer(
- AsymmetricHashing ash,
AsymmetricHashing.QuantizedVector[] vectors,
float[] scales,
float[] offsets,
+ byte[] landmarks,
long[] packedBits,
int blockSize,
- VectorFloat> query) {
+ QueryPrecompute qp) {
+
+ this.vecUtil = VectorizationProvider.getInstance().getVectorUtilSupport();
- this.vecUtil =
- io.github.jbellis.jvector.vector.VectorizationProvider
- .getInstance()
- .getVectorUtilSupport();
- this.ash = ash;
this.vectors = vectors;
this.scales = scales;
this.offsets = offsets;
+ this.landmarks = landmarks;
+
this.packedBits = packedBits;
this.blockSize = blockSize;
- this.d = ash.quantizedDim;
+ this.d = qp.d;
this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
this.maskedAdds = new float[blockSize];
- // ---- Query precompute (identical to ASHScorer) ----
- final int D = ash.originalDimension;
- final VectorFloat> mu = ash.landmarks[0];
-
- final float[] qArr = new float[D];
- final float[] muArr = new float[D];
- for (int i = 0; i < D; i++) {
- qArr[i] = query.get(i);
- muArr[i] = mu.get(i);
- }
-
- this.tildeQ = new float[d];
- final float[][] A = ash.stiefelTransform.AFloat;
- for (int r = 0; r < d; r++) {
- float s = 0f;
- float[] Arow = A[r];
- for (int j = 0; j < D; j++) {
- s += Arow[j] * (qArr[j] - muArr[j]);
- }
- tildeQ[r] = s;
- }
-
- float sum = 0f;
- for (int i = 0; i < d; i++) sum += tildeQ[i];
- this.sumTildeQ = sum;
- this.dotQMu = VectorUtil.dotProduct(query, mu);
+ this.tildeQPool = qp.tildeQPool;
+ this.tildeQScratch = new float[d];
+ this.sumTildeQ = qp.sumTildeQ;
+ this.dotQMu = qp.dotQMu;
}
+ // SIMD kernel assumes homogeneous landmark, caller handles boundaries
@Override
public void scoreRange(int start, int count, float[] out) {
+ if (count < 0) throw new IllegalArgumentException("count must be >= 0");
+ if (start < 0 || start + count > vectors.length) {
+ throw new IllegalArgumentException("Range out of bounds: [" + start + ", " + (start + count) + ")");
+ }
+ if (out.length < count) throw new IllegalArgumentException("out.length < count");
if (count == 0) return;
int ord = start;
- int end = start + count;
+ final int end = start + count;
while (ord < end) {
+ // Identify landmark for this run
+ final int c = landmarks[ord] & 0xFF;
+ final int base = c * d;
+
+ // Compute packed-block boundary info
final int blockId = ord / blockSize;
final int blockBase = blockId * blockSize;
final int laneStart = ord - blockBase;
- final int blockLen = Math.min(blockSize - laneStart, end - ord);
+ // Max we can do without crossing packed-block boundary or request end
+ final int maxLen = Math.min(blockSize - laneStart, end - ord);
+
+ // Clamp further so we do NOT cross a landmark boundary inside this packed block
+ int runLen = 1;
+ while (runLen < maxLen && ((landmarks[ord + runLen] & 0xFF) == c)) {
+ runLen++;
+ }
+
+ // Catch mixed-landmark calls for debug. TODO remove later
+ assert runLen == maxLen : "Caller passed mixed-landmark range into SIMD scorer";
+
+ final float runSumTildeQ = sumTildeQ[c];
+ final float runDotQMu = dotQMu[c];
final int blockWordBase = blockId * words * blockSize;
- // SIMD masked-add with register accumulation
vecUtil.ashMaskedAddBlockAllWords(
- tildeQ,
+ tildeQScratch,
d,
packedBits,
blockWordBase,
words,
blockSize,
laneStart,
- blockLen,
+ runLen,
maskedAdds
);
- // Finish Eq. 11 per vector
- for (int lane = 0; lane < blockLen; lane++) {
+
+ for (int lane = 0; lane < runLen; lane++) {
int idx = ord + lane;
- float scale = scales[idx];
- float offset = offsets[idx];
float m = maskedAdds[lane];
-
out[idx - start] =
- scale * (2f * m - sumTildeQ)
- + dotQMu
- + offset;
+ scales[idx] * (2f * m - runSumTildeQ)
+ + runDotQMu
+ + offsets[idx];
}
- ord += blockLen;
+ ord += runLen;
}
}
}
+ private QueryPrecompute precomputeQuery(VectorFloat> query) {
+ final int d = ash.quantizedDim;
+ final int D = ash.originalDimension;
+ final int C = ash.landmarkCount;
+
+ final float[][] A = ash.stiefelTransform.AFloat;
+
+ final var vecUtil =
+ VectorizationProvider.getInstance().getVectorUtilSupport();
+
+ // q materialized once
+ final float[] qArr = new float[D];
+ for (int i = 0; i < D; i++) {
+ qArr[i] = query.get(i);
+ }
+
+ // Aq
+ final float[] qProj = new float[d];
+ float sumQProj = 0f;
+ for (int j = 0; j < d; j++) {
+ float v = vecUtil.ashDotRow(A[j], qArr);
+ qProj[j] = v;
+ sumQProj += v;
+ }
+
+ final float[] tildeQPool = new float[C * d];
+ final float[] sumTildeQ = new float[C];
+ final float[] dotQMu = new float[C];
+
+ for (int c = 0; c < C; c++) {
+ dotQMu[c] = VectorUtil.dotProduct(query, ash.landmarks[c]);
+ sumTildeQ[c] = sumQProj - ash.landmarkProjSum[c];
+
+ final float[] muProj = ash.landmarkProj[c];
+ final int base = c * d;
+ for (int j = 0; j < d; j++) {
+ tildeQPool[base + j] = qProj[j] - muProj[j];
+ }
+ }
+
+ return new QueryPrecompute(d, C, tildeQPool, sumTildeQ, dotQMu);
+ }
+
+ static final class QueryPrecompute {
+ final int d;
+ final int C;
+ final float[] tildeQPool; // [C * d]
+ final float[] sumTildeQ; // [C]
+ final float[] dotQMu; // [C]
+
+ QueryPrecompute(int d, int C,
+ float[] tildeQPool,
+ float[] sumTildeQ,
+ float[] dotQMu) {
+ this.d = d;
+ this.C = C;
+ this.tildeQPool = tildeQPool;
+ this.sumTildeQ = sumTildeQ;
+ this.dotQMu = dotQMu;
+ }
+ }
+
/**
* For ASH, precomputed and non-precomputed scoring are currently identical.
*
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 e40bcd862..32d1f8357 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
@@ -306,4 +306,55 @@ default void ashMaskedAddBlockAllWords(
outMaskedAdd[lane] = maskedAdd;
}
}
+
+ /**
+ * Computes maskedAdd for a block slice using a pooled landmark-specific query buffer.
+ *
+ * This is identical to {@link #ashMaskedAddBlockAllWords(float[], int, long[], int, int, int, int, int, float[])}
+ * except that the query vector is read from {@code tildeQPool} at offsets {@code tildeBase + j}.
+ *
+ * tildeQPool layout:
+ *
+ * - {@code tildeQPool[tildeBase + j]} is q̃_c[j] for the caller-selected landmark c
+ *
+ *
+ * Default implementation is scalar bit-walk with register accumulation, reading from the pool.
+ */
+ default void ashMaskedAddBlockAllWordsPooled(
+ float[] tildeQPool,
+ int tildeBase, // base offset into tildeQPool for this landmark (e.g., c * d)
+ int d,
+ long[] packedBits,
+ int packedBase, // blockId * nWords * blockSize
+ int nWords,
+ int blockSize,
+ int laneStart,
+ int blockLen,
+ float[] outMaskedAdd
+ ) {
+ assert tildeBase >= 0;
+ assert laneStart >= 0;
+ assert blockLen >= 0;
+ assert laneStart + blockLen <= blockSize;
+ assert outMaskedAdd.length >= blockLen : "outMaskedAdd too small";
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ float maskedAdd = 0f; // register accumulator
+ int laneIndex = laneStart + lane;
+
+ for (int w = 0; w < nWords; w++) {
+ long word = packedBits[packedBase + (w * blockSize) + laneIndex];
+ int baseDim = w * 64;
+
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int j = baseDim + bit;
+ if (j < d) maskedAdd += tildeQPool[tildeBase + j];
+ word &= (word - 1);
+ }
+ }
+
+ outMaskedAdd[lane] = maskedAdd;
+ }
+ }
}
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 33bab6aea..d65364ea9 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
@@ -1659,6 +1659,73 @@ public void ashMaskedAddBlockAllWords(float[] tildeQ,
}
}
+ @Override
+ public void ashMaskedAddBlockAllWordsPooled(float[] tildeQPool,
+ int tildeBase,
+ int d,
+ long[] packedBits,
+ int blockWordBase,
+ int words,
+ int blockSize,
+ int laneStart,
+ int blockLen,
+ float[] outMaskedAdd) {
+
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int LANES = SPEC.length();
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ FloatVector accVec = FloatVector.zero(SPEC);
+ float scalarAcc = 0f;
+
+ int laneIndex = laneStart + lane;
+
+ for (int w = 0; w < words; w++) {
+ long wordBits = packedBits[blockWordBase + (w * blockSize) + laneIndex];
+ int baseDim = w * 64;
+
+ for (int off = 0; off < 64 && (baseDim + off) < d; off += LANES) {
+ int qBase = baseDim + off;
+
+ VectorMask inRange = SPEC.indexInRange(qBase, d);
+
+ long maskBits;
+ if (LANES == 16) {
+ maskBits = (wordBits >>> off) & 0xFFFFL;
+ } else if (LANES == 8) {
+ maskBits = (wordBits >>> off) & 0xFFL;
+ } else if (LANES == 4) {
+ maskBits = (wordBits >>> off) & 0xFL;
+ } else {
+ long m = wordBits >>> off;
+ float s = 0f;
+ for (int b = 0; b < LANES; b++) {
+ if (((m >>> b) & 1L) != 0L) {
+ int idx = qBase + b;
+ if (idx < d) s += tildeQPool[tildeBase + idx];
+ }
+ }
+ scalarAcc += s;
+ continue;
+ }
+
+ if (maskBits == 0L) continue;
+
+ VectorMask bitMask =
+ VectorMask.fromLong(SPEC, maskBits).and(inRange);
+
+ // pooled masked load: note tildeBase + qBase
+ FloatVector masked =
+ FloatVector.fromArray(SPEC, tildeQPool, tildeBase + qBase, bitMask);
+
+ accVec = accVec.add(masked);
+ }
+ }
+
+ outMaskedAdd[lane] = accVec.reduceLanes(VectorOperators.ADD) + scalarAcc;
+ }
+ }
+
private static void ashMaskedAddBlockWordLutNibble(float[] tildeQ,
int baseDim,
int d,
From d903eb604d5821de6f62926bcf47b710183e5f5f Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Mon, 5 Jan 2026 18:07:05 +0000
Subject: [PATCH 021/243] Added debug instrumentation to kmeans++ clusterer.
---
.../quantization/KMeansPlusPlusClusterer.java | 71 +++++++++++++++++--
1 file changed, 67 insertions(+), 4 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java
index 13d683327..c03acdbd3 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java
@@ -38,6 +38,13 @@
public class KMeansPlusPlusClusterer {
private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+ // Observability flags (disabled by default)
+ private static final boolean VERBOSE_KMEANS =
+ Boolean.getBoolean("jvector.kmeans.verbose");
+
+ private static final boolean VERBOSE_KMEANS_TIMING =
+ Boolean.getBoolean("jvector.kmeans.timing");
+
public static final float UNWEIGHTED = -1.0f;
// number of centroids to compute
@@ -129,18 +136,74 @@ static float computeParallelCostMultiplier(double threshold, int dimensions) {
* @return a VectorFloat of cluster centroids.
*/
public VectorFloat> cluster(int unweightedIterations, int anisotropicIterations) {
- // Always cluster unweighted first, it is significantly faster
+
+ final int n = points.length;
+ final int threshold = Math.max(1, (int) (0.01 * n));
+
+ // ------------------------------------------------------------
+ // Unweighted (isotropic) k-means phase
+ // ------------------------------------------------------------
for (int i = 0; i < unweightedIterations; i++) {
+
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
int changedCount = clusterOnceUnweighted();
- if (changedCount <= 0.01 * points.length) {
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ if (VERBOSE_KMEANS) {
+ System.out.printf(
+ "KMeans unweighted iter %d: %d/%d reassigned (%.3f%%)%n",
+ i,
+ changedCount,
+ n,
+ 100.0 * changedCount / n
+ );
+ }
+
+ if (VERBOSE_KMEANS_TIMING) {
+ System.out.printf(
+ " time = %.3f ms%n",
+ (t1 - t0) * 1e-6
+ );
+ }
+
+ if (changedCount <= threshold) {
+ if (VERBOSE_KMEANS) {
+ System.out.println("KMeans unweighted converged early");
+ }
break;
}
}
- // Optionally, refine using anisotropic clustering
+ // ------------------------------------------------------------
+ // Anisotropic refinement phase
+ // ------------------------------------------------------------
for (int i = 0; i < anisotropicIterations; i++) {
+
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
int changedCount = clusterOnceAnisotropic();
- if (changedCount <= 0.01 * points.length) {
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ if (VERBOSE_KMEANS) {
+ System.out.printf(
+ "KMeans anisotropic iter %d: %d/%d reassigned (%.3f%%)%n",
+ i,
+ changedCount,
+ n,
+ 100.0 * changedCount / n
+ );
+ }
+
+ if (VERBOSE_KMEANS_TIMING) {
+ System.out.printf(
+ " time = %.3f ms%n",
+ (t1 - t0) * 1e-6
+ );
+ }
+
+ if (changedCount <= threshold) {
+ if (VERBOSE_KMEANS) {
+ System.out.println("KMeans anisotropic converged early");
+ }
break;
}
}
From db5bfbd0eafd13a757509a75b4477b2843fb5746 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Tue, 6 Jan 2026 05:25:32 +0000
Subject: [PATCH 022/243] Bug fixes in SIMD scoreRange for copying queries into
contiguous scratch and mixed landmarks.
---
.../jvector/quantization/ASHVectors.java | 27 +++++++++++++++++--
.../jbellis/jvector/example/DistancesASH.java | 20 +++++++-------
2 files changed, 35 insertions(+), 12 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index 57ea8debf..db9579a0d 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -470,6 +470,9 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private final VectorUtilSupport vecUtil;
+ private static final java.util.concurrent.atomic.AtomicBoolean PRINTED_MIXED_ONCE =
+ new java.util.concurrent.atomic.AtomicBoolean(false);
+
BlockSimdASHBlockScorer(
AsymmetricHashing.QuantizedVector[] vectors,
float[] scales,
@@ -500,7 +503,7 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
this.dotQMu = qp.dotQMu;
}
- // SIMD kernel assumes homogeneous landmark, caller handles boundaries
+ // SIMD kernel is invoked on homogeneous landmark runs; this method splits mixed ranges into runs.
@Override
public void scoreRange(int start, int count, float[] out) {
if (count < 0) throw new IllegalArgumentException("count must be >= 0");
@@ -533,12 +536,32 @@ public void scoreRange(int start, int count, float[] out) {
}
// Catch mixed-landmark calls for debug. TODO remove later
- assert runLen == maxLen : "Caller passed mixed-landmark range into SIMD scorer";
+ if (runLen != maxLen && PRINTED_MIXED_ONCE.compareAndSet(false, true)) {
+ int nextLm = (ord + runLen < end) ? (landmarks[ord + runLen] & 0xFF) : -1;
+ System.out.println(
+ "\t[ASH SIMD debug] mixed-landmark boundary inside packed-block slice: " +
+ "start=" + start +
+ " count=" + count +
+ " ord=" + ord +
+ " blockId=" + blockId +
+ " laneStart=" + laneStart +
+ " runLen=" + runLen +
+ " maxLen=" + maxLen +
+ " c=" + c +
+ " nextC=" + nextLm
+ );
+ }
+
+ assert runLen >= 1 && runLen <= maxLen;
+ assert laneStart >= 0 && laneStart + runLen <= blockSize;
final float runSumTildeQ = sumTildeQ[c];
final float runDotQMu = dotQMu[c];
final int blockWordBase = blockId * words * blockSize;
+ // Load q̃_c into contiguous scratch for the unpooled kernel
+ System.arraycopy(tildeQPool, base, tildeQScratch, 0, d);
+
vecUtil.ashMaskedAddBlockAllWords(
tildeQScratch,
d,
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 37f0a20cf..3e5f841c7 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -64,14 +64,14 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
final int HEADER_BITS = 72;
// Block sizes to benchmark
- final int[] BLOCK_SIZES = {64}; // 16, 32, and/or 64
+ final int[] BLOCK_SIZES = {16, 32, 64}; // 16, 32, and/or 64
// How many ASH landmarks to use, C = [1, 64]
- final int landmarkCount = 64;
+ final int landmarkCount = 1;
// Define the benchmark size
- int maxQueries = 10_000;
- int maxVectors = 10_000_000;
+ int maxQueries = 1_000;
+ int maxVectors = 1_000_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -509,11 +509,11 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
-// runCohere100k();
-// runADA();
-// runOpenai1536();
-// runOpenai3072();
- runCap6m();
- runCohere10m();
+ runCohere100k();
+ runADA();
+ runOpenai1536();
+ runOpenai3072();
+// runCap6m();
+// runCohere10m();
}
}
From 2adbce00f73cde98328f079aa01e52675d2f489d Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Wed, 7 Jan 2026 02:37:17 +0000
Subject: [PATCH 023/243] Encoded vectors are now presorted by landmark for
faster full dataset scoring. Also made a slight change to limit copying
query info from query pool (only at landmark run boundaries) during block
scoring.
---
.../jvector/quantization/ASHVectors.java | 206 ++++++++++++++++--
.../jbellis/jvector/example/DistancesASH.java | 59 ++++-
2 files changed, 244 insertions(+), 21 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index db9579a0d..488dc43a2 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -112,6 +112,36 @@ public ASHVectors(AsymmetricHashing ash,
}
+ // Private constructor that trusts prebuilt arrays (for vector landmark sorting)
+ private ASHVectors(AsymmetricHashing ash,
+ AsymmetricHashing.QuantizedVector[] compressedVectors,
+ float[] scales,
+ float[] offsets,
+ byte[] landmarks) {
+ this.ash = ash;
+ this.compressedVectors = compressedVectors;
+ this.scorer = new ASHScorer(ash);
+
+ this.d = ash.quantizedDim;
+ this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
+
+ // Defensive sanity checks (cheap)
+ if (scales.length != compressedVectors.length ||
+ offsets.length != compressedVectors.length ||
+ landmarks.length != compressedVectors.length) {
+ throw new IllegalArgumentException("Header arrays must match compressedVectors length");
+ }
+
+ this.scales = scales;
+ this.offsets = offsets;
+ this.landmarks = landmarks;
+
+ // packedBits cache starts empty (as usual)
+ this.packedBlockSize = -1;
+ this.packedBits = null;
+ this.packedBlockCount = 0;
+ }
+
@Override
public int count() {
return compressedVectors.length;
@@ -468,6 +498,9 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private final float[] sumTildeQ; // qp.sumTildeQ [C]
private final float[] dotQMu; // qp.dotQMu [C]
+ // Cache which landmark's q̃ is currently loaded into tildeQScratch
+ private int lastC = -1;
+
private final VectorUtilSupport vecUtil;
private static final java.util.concurrent.atomic.AtomicBoolean PRINTED_MIXED_ONCE =
@@ -536,21 +569,21 @@ public void scoreRange(int start, int count, float[] out) {
}
// Catch mixed-landmark calls for debug. TODO remove later
- if (runLen != maxLen && PRINTED_MIXED_ONCE.compareAndSet(false, true)) {
- int nextLm = (ord + runLen < end) ? (landmarks[ord + runLen] & 0xFF) : -1;
- System.out.println(
- "\t[ASH SIMD debug] mixed-landmark boundary inside packed-block slice: " +
- "start=" + start +
- " count=" + count +
- " ord=" + ord +
- " blockId=" + blockId +
- " laneStart=" + laneStart +
- " runLen=" + runLen +
- " maxLen=" + maxLen +
- " c=" + c +
- " nextC=" + nextLm
- );
- }
+// if (runLen != maxLen && PRINTED_MIXED_ONCE.compareAndSet(false, true)) {
+// int nextLm = (ord + runLen < end) ? (landmarks[ord + runLen] & 0xFF) : -1;
+// System.out.println(
+// "\t[ASH SIMD debug] mixed-landmark boundary inside packed-block slice: " +
+// "start=" + start +
+// " count=" + count +
+// " ord=" + ord +
+// " blockId=" + blockId +
+// " laneStart=" + laneStart +
+// " runLen=" + runLen +
+// " maxLen=" + maxLen +
+// " c=" + c +
+// " nextC=" + nextLm
+// );
+// }
assert runLen >= 1 && runLen <= maxLen;
assert laneStart >= 0 && laneStart + runLen <= blockSize;
@@ -560,7 +593,13 @@ public void scoreRange(int start, int count, float[] out) {
final int blockWordBase = blockId * words * blockSize;
// Load q̃_c into contiguous scratch for the unpooled kernel
+ if (c != lastC) {
+ System.arraycopy(tildeQPool, base, tildeQScratch, 0, d);
+ lastC = c;
+ }
+
System.arraycopy(tildeQPool, base, tildeQScratch, 0, d);
+ lastC = c;
vecUtil.ashMaskedAddBlockAllWords(
tildeQScratch,
@@ -575,6 +614,8 @@ public void scoreRange(int start, int count, float[] out) {
);
+
+
for (int lane = 0; lane < runLen; lane++) {
int idx = ord + lane;
float m = maskedAdds[lane];
@@ -651,6 +692,141 @@ static final class QueryPrecompute {
}
}
+ public static final class LandmarkOrder {
+ public final ASHVectors vectors;
+ public final int[] newToOld;
+ public final int[] oldToNew;
+ public final int[] landmarkOffsets; // length C+1, last == n
+
+ LandmarkOrder(ASHVectors vectors, int[] newToOld, int[] oldToNew, int[] landmarkOffsets) {
+ this.vectors = vectors;
+ this.newToOld = newToOld;
+ this.oldToNew = oldToNew;
+ this.landmarkOffsets = landmarkOffsets;
+ }
+ }
+
+ /**
+ * Stable, linear-time reorder by landmark id (C <= 64).
+ * Fills reordered header arrays in the same pass (no second constructor scan).
+ */
+ public LandmarkOrder reorderByLandmarkFast() {
+ final int n = compressedVectors.length;
+ final int C = ash.landmarkCount;
+
+ // Count per landmark
+ final int[] counts = new int[C];
+ for (int i = 0; i < n; i++) {
+ int c = this.landmarks[i] & 0xFF;
+ if (c >= C) {
+ throw new IllegalStateException("Invalid landmark id " + c + " at ordinal " + i + " (C=" + C + ")");
+ }
+ counts[c]++;
+ }
+
+ // Prefix sums: landmarkOffsets[c] is start of landmark c, length C+1
+ final int[] landmarkOffsets = new int[C + 1];
+ landmarkOffsets[0] = 0;
+ for (int c = 0; c < C; c++) {
+ landmarkOffsets[c + 1] = landmarkOffsets[c] + counts[c];
+ }
+ if (landmarkOffsets[C] != n) {
+ throw new IllegalStateException("landmarkOffsets[C] != n: " + landmarkOffsets[C] + " != " + n);
+ }
+
+ final int[] writePos = landmarkOffsets.clone();
+
+ final AsymmetricHashing.QuantizedVector[] newVecs = new AsymmetricHashing.QuantizedVector[n];
+ final float[] newScales = new float[n];
+ final float[] newOffsets = new float[n];
+ final byte[] newLandmarks = new byte[n];
+ final int[] newToOld = new int[n];
+
+ // Preserve an ordinal map so we can go back and forth
+ for (int oldOrd = 0; oldOrd < n; oldOrd++) {
+ final int c = this.landmarks[oldOrd] & 0xFF;
+ final int newOrd = writePos[c]++;
+
+ newVecs[newOrd] = this.compressedVectors[oldOrd];
+ newScales[newOrd] = this.scales[oldOrd];
+ newOffsets[newOrd] = this.offsets[oldOrd];
+ newLandmarks[newOrd] = this.landmarks[oldOrd];
+
+ newToOld[newOrd] = oldOrd;
+ }
+
+ final int[] oldToNew = new int[n];
+ for (int newOrd = 0; newOrd < n; newOrd++) {
+ oldToNew[newToOld[newOrd]] = newOrd;
+ }
+
+ ASHVectors reordered = new ASHVectors(ash, newVecs, newScales, newOffsets, newLandmarks);
+ return new LandmarkOrder(reordered, newToOld, oldToNew, landmarkOffsets);
+ }
+
+ public static final class LandmarkRunStats {
+ public final int n;
+ public final int runCount;
+ public final double avgRunLength;
+ public final int maxRunLength;
+ public final double stddevRunLength;
+
+ LandmarkRunStats(int n, int runCount, double avgRunLength, int maxRunLength, double stddevRunLength) {
+ this.n = n;
+ this.runCount = runCount;
+ this.avgRunLength = avgRunLength;
+ this.maxRunLength = maxRunLength;
+ this.stddevRunLength = stddevRunLength;
+ }
+
+ @Override
+ public String toString() {
+ return String.format(
+ java.util.Locale.ROOT,
+ "runs=%d, avgRun=%.2f, maxRun=%d, stddev=%.2f (n=%d)",
+ runCount, avgRunLength, maxRunLength, stddevRunLength, n
+ );
+ }
+ }
+
+ public LandmarkRunStats landmarkRunStats() {
+ return computeLandmarkRunStats(this.landmarks);
+ }
+
+ public static LandmarkRunStats computeLandmarkRunStats(byte[] landmarks) {
+ final int n = landmarks.length;
+ if (n == 0) {
+ return new LandmarkRunStats(0, 0, 0.0, 0, 0.0);
+ }
+
+ int runCount = 0;
+ long sum = 0;
+ long sumSq = 0;
+ int max = 0;
+
+ int i = 0;
+ while (i < n) {
+ final byte c = landmarks[i];
+ int j = i + 1;
+ while (j < n && landmarks[j] == c) j++;
+ final int len = j - i;
+
+ runCount++;
+ sum += len;
+ sumSq += (long) len * (long) len;
+ if (len > max) max = len;
+
+ i = j;
+ }
+
+ final double avg = (double) sum / (double) runCount;
+ final double meanSq = (double) sumSq / (double) runCount;
+ final double var = Math.max(0.0, meanSq - avg * avg);
+ final double stddev = Math.sqrt(var);
+
+ return new LandmarkRunStats(n, runCount, avg, max, stddev);
+ }
+
/**
* For ASH, precomputed and non-precomputed scoring are currently identical.
*
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 3e5f841c7..49baff0c7 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -52,7 +52,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "true"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
final boolean RUN_SCALAR_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
@@ -171,6 +171,49 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
ASHVectors ashVectors = (ASHVectors) ashVecs;
VectorFloat> q0 = finalQueries.get(0);
+ // ------------------------------------------------------------------
+ // Landmark reorder (debug + timing)
+ // ------------------------------------------------------------------
+ final boolean REORDER_BY_LANDMARK =
+ Boolean.parseBoolean(System.getProperty("jvector.ash.reorderByLandmark", "true"));
+
+ final ASHVectors ashVectorsFinal;
+ final CompressedVectors ashVecsFinal;
+ final int[] newToOldFinal;
+
+ if (REORDER_BY_LANDMARK) {
+ long stat0 = System.nanoTime();
+ var beforeStats = ashVectors.landmarkRunStats();
+ long stat1 = System.nanoTime();
+
+ System.out.println("\tLandmark run stats BEFORE reorder: " + beforeStats
+ + " (computed in " + (stat1 - stat0) / 1e9 + " s)");
+
+ long t0 = System.nanoTime();
+ ASHVectors.LandmarkOrder order = ashVectors.reorderByLandmarkFast();
+ long t1 = System.nanoTime();
+
+ ASHVectors reordered = order.vectors;
+
+ long stat2 = System.nanoTime();
+ var afterStats = reordered.landmarkRunStats();
+ long stat3 = System.nanoTime();
+
+ System.out.println("\tLandmark reorder took " + (t1 - t0) / 1e9 + " s");
+ System.out.println("\tLandmark run stats AFTER reorder: " + afterStats
+ + " (computed in " + (stat3 - stat2) / 1e9 + " s)");
+ System.out.println("--");
+
+ ashVectorsFinal = reordered;
+ ashVecsFinal = reordered; // CompressedVectors view used by scoreFunctionFor
+ newToOldFinal = order.newToOld; // map reordered ordinal -> original base ordinal
+ } else {
+ System.out.println("\tLandmark reorder disabled (-Djvector.ash.reorderByLandmark=false)");
+ ashVectorsFinal = ashVectors;
+ ashVecsFinal = ashVectors;
+ newToOldFinal = null;
+ }
+
if(RUN_SANITY_CHECK) {
ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
@@ -226,16 +269,19 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (int i = s; i < e; i++) {
VectorFloat> q = finalQueries.get(i);
ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ ashVecsFinal.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
for (int j = 0; j < finalVectorCount; j++) {
- VectorFloat> v = finalVectors.get(j);
+ final int baseOrd = (newToOldFinal == null) ? j : newToOldFinal[j];
+ VectorFloat> v = finalVectors.get(baseOrd);
+
float trueDot = VectorUtil.dotProduct(q, v);
float approxDot = f.similarityTo(j);
localError += Math.abs(approxDot - trueDot);
localCount++;
}
+
}
return new double[]{localError, localCount};
}));
@@ -348,7 +394,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Print scorer implementation once per blockSize (diagnostic)
{
ASHBlockScorer scorer =
- ashVectors.blockScorerFor(
+ ashVectorsFinal.blockScorerFor(
finalQueries.get(0),
VectorSimilarityFunction.DOT_PRODUCT,
blockSize
@@ -370,7 +416,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (int qi = s; qi < e; qi++) {
ASHBlockScorer scorer =
- ashVectors.blockScorerFor(
+ ashVectorsFinal.blockScorerFor(
finalQueries.get(qi),
VectorSimilarityFunction.DOT_PRODUCT,
blockSize
@@ -388,10 +434,11 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Tail uses scalar per-vector scorer (correctness-preserving)
if (j < finalVectorCount) {
ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(
+ ashVecsFinal.scoreFunctionFor(
finalQueries.get(qi),
VectorSimilarityFunction.DOT_PRODUCT
);
+
for (; j < finalVectorCount; j++) {
localSum += f.similarityTo(j);
}
From 70eb02e9a149f0865fde6afd9747cbec6e543e44 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Wed, 7 Jan 2026 02:37:17 +0000
Subject: [PATCH 024/243] Encoded vectors are now presorted by landmark for
faster full dataset scoring. Also made a slight change to limit copying
query info from query pool (only at landmark run boundaries) during block
scoring.
---
.../jvector/quantization/ASHVectors.java | 205 ++++++++++++++++--
.../jbellis/jvector/example/DistancesASH.java | 59 ++++-
2 files changed, 242 insertions(+), 22 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index db9579a0d..0ae65178a 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -112,6 +112,36 @@ public ASHVectors(AsymmetricHashing ash,
}
+ // Private constructor that trusts prebuilt arrays (for vector landmark sorting)
+ private ASHVectors(AsymmetricHashing ash,
+ AsymmetricHashing.QuantizedVector[] compressedVectors,
+ float[] scales,
+ float[] offsets,
+ byte[] landmarks) {
+ this.ash = ash;
+ this.compressedVectors = compressedVectors;
+ this.scorer = new ASHScorer(ash);
+
+ this.d = ash.quantizedDim;
+ this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
+
+ // Defensive sanity checks (cheap)
+ if (scales.length != compressedVectors.length ||
+ offsets.length != compressedVectors.length ||
+ landmarks.length != compressedVectors.length) {
+ throw new IllegalArgumentException("Header arrays must match compressedVectors length");
+ }
+
+ this.scales = scales;
+ this.offsets = offsets;
+ this.landmarks = landmarks;
+
+ // packedBits cache starts empty (as usual)
+ this.packedBlockSize = -1;
+ this.packedBits = null;
+ this.packedBlockCount = 0;
+ }
+
@Override
public int count() {
return compressedVectors.length;
@@ -468,6 +498,9 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private final float[] sumTildeQ; // qp.sumTildeQ [C]
private final float[] dotQMu; // qp.dotQMu [C]
+ // Cache which landmark's q̃ is currently loaded into tildeQScratch
+ private int lastC = -1;
+
private final VectorUtilSupport vecUtil;
private static final java.util.concurrent.atomic.AtomicBoolean PRINTED_MIXED_ONCE =
@@ -536,21 +569,21 @@ public void scoreRange(int start, int count, float[] out) {
}
// Catch mixed-landmark calls for debug. TODO remove later
- if (runLen != maxLen && PRINTED_MIXED_ONCE.compareAndSet(false, true)) {
- int nextLm = (ord + runLen < end) ? (landmarks[ord + runLen] & 0xFF) : -1;
- System.out.println(
- "\t[ASH SIMD debug] mixed-landmark boundary inside packed-block slice: " +
- "start=" + start +
- " count=" + count +
- " ord=" + ord +
- " blockId=" + blockId +
- " laneStart=" + laneStart +
- " runLen=" + runLen +
- " maxLen=" + maxLen +
- " c=" + c +
- " nextC=" + nextLm
- );
- }
+// if (runLen != maxLen && PRINTED_MIXED_ONCE.compareAndSet(false, true)) {
+// int nextLm = (ord + runLen < end) ? (landmarks[ord + runLen] & 0xFF) : -1;
+// System.out.println(
+// "\t[ASH SIMD debug] mixed-landmark boundary inside packed-block slice: " +
+// "start=" + start +
+// " count=" + count +
+// " ord=" + ord +
+// " blockId=" + blockId +
+// " laneStart=" + laneStart +
+// " runLen=" + runLen +
+// " maxLen=" + maxLen +
+// " c=" + c +
+// " nextC=" + nextLm
+// );
+// }
assert runLen >= 1 && runLen <= maxLen;
assert laneStart >= 0 && laneStart + runLen <= blockSize;
@@ -560,7 +593,10 @@ public void scoreRange(int start, int count, float[] out) {
final int blockWordBase = blockId * words * blockSize;
// Load q̃_c into contiguous scratch for the unpooled kernel
- System.arraycopy(tildeQPool, base, tildeQScratch, 0, d);
+ if (c != lastC) {
+ System.arraycopy(tildeQPool, base, tildeQScratch, 0, d);
+ lastC = c;
+ }
vecUtil.ashMaskedAddBlockAllWords(
tildeQScratch,
@@ -575,6 +611,8 @@ public void scoreRange(int start, int count, float[] out) {
);
+
+
for (int lane = 0; lane < runLen; lane++) {
int idx = ord + lane;
float m = maskedAdds[lane];
@@ -651,6 +689,141 @@ static final class QueryPrecompute {
}
}
+ public static final class LandmarkOrder {
+ public final ASHVectors vectors;
+ public final int[] newToOld;
+ public final int[] oldToNew;
+ public final int[] landmarkOffsets; // length C+1, last == n
+
+ LandmarkOrder(ASHVectors vectors, int[] newToOld, int[] oldToNew, int[] landmarkOffsets) {
+ this.vectors = vectors;
+ this.newToOld = newToOld;
+ this.oldToNew = oldToNew;
+ this.landmarkOffsets = landmarkOffsets;
+ }
+ }
+
+ /**
+ * Stable, linear-time reorder by landmark id (C <= 64).
+ * Fills reordered header arrays in the same pass (no second constructor scan).
+ */
+ public LandmarkOrder reorderByLandmarkFast() {
+ final int n = compressedVectors.length;
+ final int C = ash.landmarkCount;
+
+ // Count per landmark
+ final int[] counts = new int[C];
+ for (int i = 0; i < n; i++) {
+ int c = this.landmarks[i] & 0xFF;
+ if (c >= C) {
+ throw new IllegalStateException("Invalid landmark id " + c + " at ordinal " + i + " (C=" + C + ")");
+ }
+ counts[c]++;
+ }
+
+ // Prefix sums: landmarkOffsets[c] is start of landmark c, length C+1
+ final int[] landmarkOffsets = new int[C + 1];
+ landmarkOffsets[0] = 0;
+ for (int c = 0; c < C; c++) {
+ landmarkOffsets[c + 1] = landmarkOffsets[c] + counts[c];
+ }
+ if (landmarkOffsets[C] != n) {
+ throw new IllegalStateException("landmarkOffsets[C] != n: " + landmarkOffsets[C] + " != " + n);
+ }
+
+ final int[] writePos = landmarkOffsets.clone();
+
+ final AsymmetricHashing.QuantizedVector[] newVecs = new AsymmetricHashing.QuantizedVector[n];
+ final float[] newScales = new float[n];
+ final float[] newOffsets = new float[n];
+ final byte[] newLandmarks = new byte[n];
+ final int[] newToOld = new int[n];
+
+ // Preserve an ordinal map so we can go back and forth
+ for (int oldOrd = 0; oldOrd < n; oldOrd++) {
+ final int c = this.landmarks[oldOrd] & 0xFF;
+ final int newOrd = writePos[c]++;
+
+ newVecs[newOrd] = this.compressedVectors[oldOrd];
+ newScales[newOrd] = this.scales[oldOrd];
+ newOffsets[newOrd] = this.offsets[oldOrd];
+ newLandmarks[newOrd] = this.landmarks[oldOrd];
+
+ newToOld[newOrd] = oldOrd;
+ }
+
+ final int[] oldToNew = new int[n];
+ for (int newOrd = 0; newOrd < n; newOrd++) {
+ oldToNew[newToOld[newOrd]] = newOrd;
+ }
+
+ ASHVectors reordered = new ASHVectors(ash, newVecs, newScales, newOffsets, newLandmarks);
+ return new LandmarkOrder(reordered, newToOld, oldToNew, landmarkOffsets);
+ }
+
+ public static final class LandmarkRunStats {
+ public final int n;
+ public final int runCount;
+ public final double avgRunLength;
+ public final int maxRunLength;
+ public final double stddevRunLength;
+
+ LandmarkRunStats(int n, int runCount, double avgRunLength, int maxRunLength, double stddevRunLength) {
+ this.n = n;
+ this.runCount = runCount;
+ this.avgRunLength = avgRunLength;
+ this.maxRunLength = maxRunLength;
+ this.stddevRunLength = stddevRunLength;
+ }
+
+ @Override
+ public String toString() {
+ return String.format(
+ java.util.Locale.ROOT,
+ "runs=%d, avgRun=%.2f, maxRun=%d, stddev=%.2f (n=%d)",
+ runCount, avgRunLength, maxRunLength, stddevRunLength, n
+ );
+ }
+ }
+
+ public LandmarkRunStats landmarkRunStats() {
+ return computeLandmarkRunStats(this.landmarks);
+ }
+
+ public static LandmarkRunStats computeLandmarkRunStats(byte[] landmarks) {
+ final int n = landmarks.length;
+ if (n == 0) {
+ return new LandmarkRunStats(0, 0, 0.0, 0, 0.0);
+ }
+
+ int runCount = 0;
+ long sum = 0;
+ long sumSq = 0;
+ int max = 0;
+
+ int i = 0;
+ while (i < n) {
+ final byte c = landmarks[i];
+ int j = i + 1;
+ while (j < n && landmarks[j] == c) j++;
+ final int len = j - i;
+
+ runCount++;
+ sum += len;
+ sumSq += (long) len * (long) len;
+ if (len > max) max = len;
+
+ i = j;
+ }
+
+ final double avg = (double) sum / (double) runCount;
+ final double meanSq = (double) sumSq / (double) runCount;
+ final double var = Math.max(0.0, meanSq - avg * avg);
+ final double stddev = Math.sqrt(var);
+
+ return new LandmarkRunStats(n, runCount, avg, max, stddev);
+ }
+
/**
* For ASH, precomputed and non-precomputed scoring are currently identical.
*
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 3e5f841c7..49baff0c7 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -52,7 +52,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "true"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
final boolean RUN_SCALAR_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
@@ -171,6 +171,49 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
ASHVectors ashVectors = (ASHVectors) ashVecs;
VectorFloat> q0 = finalQueries.get(0);
+ // ------------------------------------------------------------------
+ // Landmark reorder (debug + timing)
+ // ------------------------------------------------------------------
+ final boolean REORDER_BY_LANDMARK =
+ Boolean.parseBoolean(System.getProperty("jvector.ash.reorderByLandmark", "true"));
+
+ final ASHVectors ashVectorsFinal;
+ final CompressedVectors ashVecsFinal;
+ final int[] newToOldFinal;
+
+ if (REORDER_BY_LANDMARK) {
+ long stat0 = System.nanoTime();
+ var beforeStats = ashVectors.landmarkRunStats();
+ long stat1 = System.nanoTime();
+
+ System.out.println("\tLandmark run stats BEFORE reorder: " + beforeStats
+ + " (computed in " + (stat1 - stat0) / 1e9 + " s)");
+
+ long t0 = System.nanoTime();
+ ASHVectors.LandmarkOrder order = ashVectors.reorderByLandmarkFast();
+ long t1 = System.nanoTime();
+
+ ASHVectors reordered = order.vectors;
+
+ long stat2 = System.nanoTime();
+ var afterStats = reordered.landmarkRunStats();
+ long stat3 = System.nanoTime();
+
+ System.out.println("\tLandmark reorder took " + (t1 - t0) / 1e9 + " s");
+ System.out.println("\tLandmark run stats AFTER reorder: " + afterStats
+ + " (computed in " + (stat3 - stat2) / 1e9 + " s)");
+ System.out.println("--");
+
+ ashVectorsFinal = reordered;
+ ashVecsFinal = reordered; // CompressedVectors view used by scoreFunctionFor
+ newToOldFinal = order.newToOld; // map reordered ordinal -> original base ordinal
+ } else {
+ System.out.println("\tLandmark reorder disabled (-Djvector.ash.reorderByLandmark=false)");
+ ashVectorsFinal = ashVectors;
+ ashVecsFinal = ashVectors;
+ newToOldFinal = null;
+ }
+
if(RUN_SANITY_CHECK) {
ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
@@ -226,16 +269,19 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (int i = s; i < e; i++) {
VectorFloat> q = finalQueries.get(i);
ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ ashVecsFinal.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
for (int j = 0; j < finalVectorCount; j++) {
- VectorFloat> v = finalVectors.get(j);
+ final int baseOrd = (newToOldFinal == null) ? j : newToOldFinal[j];
+ VectorFloat> v = finalVectors.get(baseOrd);
+
float trueDot = VectorUtil.dotProduct(q, v);
float approxDot = f.similarityTo(j);
localError += Math.abs(approxDot - trueDot);
localCount++;
}
+
}
return new double[]{localError, localCount};
}));
@@ -348,7 +394,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Print scorer implementation once per blockSize (diagnostic)
{
ASHBlockScorer scorer =
- ashVectors.blockScorerFor(
+ ashVectorsFinal.blockScorerFor(
finalQueries.get(0),
VectorSimilarityFunction.DOT_PRODUCT,
blockSize
@@ -370,7 +416,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (int qi = s; qi < e; qi++) {
ASHBlockScorer scorer =
- ashVectors.blockScorerFor(
+ ashVectorsFinal.blockScorerFor(
finalQueries.get(qi),
VectorSimilarityFunction.DOT_PRODUCT,
blockSize
@@ -388,10 +434,11 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
// Tail uses scalar per-vector scorer (correctness-preserving)
if (j < finalVectorCount) {
ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(
+ ashVecsFinal.scoreFunctionFor(
finalQueries.get(qi),
VectorSimilarityFunction.DOT_PRODUCT
);
+
for (; j < finalVectorCount; j++) {
localSum += f.similarityTo(j);
}
From 21ca46e7e30b1c9518ae98119c6692da72ab84ab Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Thu, 8 Jan 2026 04:57:06 +0000
Subject: [PATCH 025/243] Initial implementation of LAPACK-based ITQ trainer
and slow Landing trainer. Netlib dependencies added.
---
jvector-base/pom.xml | 22 +-
.../jvector/quantization/ASHVectors.java | 1 +
.../quantization/AsymmetricHashing.java | 545 +++++++++++++++++-
.../jbellis/jvector/example/DistancesASH.java | 45 +-
pom.xml | 28 +-
5 files changed, 627 insertions(+), 14 deletions(-)
diff --git a/jvector-base/pom.xml b/jvector-base/pom.xml
index c5405dc9d..e3c12502f 100644
--- a/jvector-base/pom.xml
+++ b/jvector-base/pom.xml
@@ -1,5 +1,5 @@
-
@@ -12,6 +12,26 @@
jvector-base
Base
+
+
+ com.github.fommil.netlib
+ core
+
+
+ dev.ludovic.netlib
+ blas
+
+
+ dev.ludovic.netlib
+ lapack
+
+
+ dev.ludovic.netlib
+ lapack
+ 3.0.3
+
+
+
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index 0ae65178a..10e66a487 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -617,6 +617,7 @@ public void scoreRange(int start, int count, float[] out) {
int idx = ord + lane;
float m = maskedAdds[lane];
out[idx - start] =
+ // scale the 0/1 stored bits to +/-1 for processing.
scales[idx] * (2f * m - runSumTildeQ)
+ runDotQMu
+ offsets[idx];
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 5f253d4f8..ddd02da3c 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -28,6 +28,8 @@
import io.github.jbellis.jvector.vector.types.VectorFloat;
import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+import dev.ludovic.netlib.blas.BLAS;
+import dev.ludovic.netlib.lapack.LAPACK;
import java.io.DataOutput;
import java.io.IOException;
import java.util.Arrays;
@@ -37,6 +39,9 @@
import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.QRDecomposition;
import org.apache.commons.math3.linear.RealMatrix;
+import org.apache.commons.math3.linear.Array2DRowRealMatrix;
+import org.apache.commons.math3.linear.SingularValueDecomposition;
+import org.netlib.util.intW;
import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED;
@@ -49,6 +54,26 @@ public class AsymmetricHashing implements VectorCompressor − ||μ||_2^2
@@ -280,13 +305,21 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
}
}
+ // Keep determinism consistent within ASH. TODO define seed constant or remove.
+ final Random rng = new Random(42);
+
StiefelTransform stiefelTransform;
if (optimizer == RANDOM) {
- stiefelTransform = runWithoutTraining(originalDim, quantizedDim, new Random(42));
- } else if (optimizer == ITQ) {
- throw new IllegalArgumentException("ITQ optimizer not implemented yet");
- } else if (optimizer == LANDING) {
- throw new IllegalArgumentException("LANDING optimizer not implemented yet");
+ stiefelTransform = runWithoutTraining(originalDim, quantizedDim, rng);
+ } else if (optimizer == ITQ || optimizer == LANDING) {
+ // Build training inputs from a random sample of points.
+ TrainingData td = prepareTrainingData(points, landmarks, optimizer, originalDim, rng);
+
+ if (optimizer == ITQ) {
+ stiefelTransform = runItqTrainer(td.xTrainNorm, quantizedDim, rng, TRAINING_ITERS);
+ } else {
+ stiefelTransform = runLandingTrainer(td.xTrainNorm, quantizedDim, rng, TRAINING_ITERS);
+ }
} else {
throw new IllegalArgumentException("Unknown optimizer " + optimizer);
}
@@ -642,6 +675,10 @@ public void quantizeBody(VectorFloat> vector,
}
public static final class ItqBinaryQuantizer implements BinaryQuantizer {
+ // Same binarization math for all optimizers; only the learned transform differs.
+ // TODO consolidate quantizeBody across learners/random
+ private final RandomBinaryQuantizer delegate = new RandomBinaryQuantizer();
+
@Override
public void quantizeBody(VectorFloat> vector,
VectorFloat> mu,
@@ -649,11 +686,14 @@ public void quantizeBody(VectorFloat> vector,
int quantizedDim,
StiefelTransform stiefel,
long[] outWords) {
- throw new UnsupportedOperationException("ITQ not implemented");
+ delegate.quantizeBody(vector, mu, residualNorm, quantizedDim, stiefel, outWords);
}
}
public static final class LandingBinaryQuantizer implements BinaryQuantizer {
+ // Same binarization math for all optimizers; only the learned transform differs.
+ private final RandomBinaryQuantizer delegate = new RandomBinaryQuantizer();
+
@Override
public void quantizeBody(VectorFloat> vector,
VectorFloat> mu,
@@ -661,7 +701,7 @@ public void quantizeBody(VectorFloat> vector,
int quantizedDim,
StiefelTransform stiefel,
long[] outWords) {
- throw new UnsupportedOperationException("LANDING not implemented");
+ delegate.quantizeBody(vector, mu, residualNorm, quantizedDim, stiefel, outWords);
}
}
@@ -839,7 +879,6 @@ public int hashCode() {
result = 31 * result + Arrays.hashCode(binaryVector);
return result;
}
-
}
/**
@@ -1149,6 +1188,491 @@ public void encodeTo(VectorFloat> vector, QuantizedVector dest) {
quantizeVector(vector, residualNorm, landmark, dest);
}
+ // ---------------------------------------------------------------------
+ // Training of projection matrix (ITQ + LANDING algorithms)
+ // ---------------------------------------------------------------------
+
+ /**
+ * Training inputs derived from a random sample of the dataset:
+ * xTrainNorm holds normalized centered residuals.
+ *
+ * NOTE: Vectors are read as floats; training matrices are accumulated/stored in double precision.
+ */
+ private static final class TrainingData {
+ final double[][] xTrainNorm; // [N_train][D]
+
+ TrainingData(double[][] xTrainNorm) {
+ this.xTrainNorm = xTrainNorm;
+ }
+ }
+
+ /**
+ * Builds normalized centered residuals for training:
+ * - Choose N_train = min(N, cap) (cap differs by optimizer).
+ * - Sample ordinals uniformly at random (deterministic RNG).
+ * - For each sampled ordinal i:
+ * c = argmin ||x_i - mu_c||^2 (same rule as encoding)
+ * r = x_i - mu_c
+ * r_hat = r / ||r||
+ * store r_hat into xTrainNorm
+ */
+ private static TrainingData prepareTrainingData(
+ VectorFloat>[] points,
+ VectorFloat>[] landmarks,
+ int optimizer,
+ int originalDim,
+ Random rng
+ ) {
+ final int N = points.length;
+ final int cap = (optimizer == LANDING) ? LANDING_N_TRAIN : ITQ_N_TRAIN;
+ final int nTrain = Math.min(N, cap);
+
+ final int[] sample = reservoirSampleOrdinals(N, nTrain, rng);
+ Arrays.sort(sample); // deterministic order for reproducibility
+
+ final double[][] xTrainNorm = new double[nTrain][originalDim];
+
+ int row = 0;
+ int samplePos = 0;
+
+ for (int i = 0; i < N && samplePos < nTrain; i++) {
+ if (sample[samplePos] != i) {
+ continue;
+ }
+ samplePos++;
+
+ final VectorFloat> x = points[i];
+
+ // Nearest landmark assignment (same semantics as encodeTo).
+ byte best = 0;
+ float bestDist = Float.POSITIVE_INFINITY;
+ for (int c = 0; c < landmarks.length; c++) {
+ float dist = VectorUtil.squareL2Distance(x, landmarks[c]);
+ if (dist < bestDist) {
+ bestDist = dist;
+ best = (byte) c;
+ }
+ }
+
+ final VectorFloat> mu = landmarks[best & 0xFF];
+
+ // Center and normalize into double precision row.
+ final double[] out = xTrainNorm[row++];
+
+ double ss = 0.0;
+ for (int d = 0; d < originalDim; d++) {
+ double v = (double) x.get(d) - (double) mu.get(d);
+ out[d] = v;
+ ss += v * v;
+ }
+
+ // Normalize residuals. NOTE: Does not handle zero vectors!
+ double inv = 1.0 / Math.sqrt(ss);
+ for (int d = 0; d < originalDim; d++) {
+ out[d] *= inv;
+ }
+ }
+
+ if (row != nTrain) {
+ throw new IllegalStateException("Training sample fill mismatch: got " + row + ", expected " + nTrain);
+ }
+
+ return new TrainingData(xTrainNorm);
+ }
+
+ /**
+ * Reservoir sample k ordinals uniformly from [0, N).
+ * Deterministic given rng.
+ */
+ private static int[] reservoirSampleOrdinals(int N, int k, Random rng) {
+ if (k < 0 || k > N) throw new IllegalArgumentException("k out of range: " + k + " for N=" + N);
+ int[] r = new int[k];
+ for (int i = 0; i < k; i++) r[i] = i;
+ for (int i = k; i < N; i++) {
+ int j = rng.nextInt(i + 1);
+ if (j < k) r[j] = i;
+ }
+ return r;
+ }
+
+ /**
+ * ITQ trainer:
+ * - PCA basis from SVD(X_train_norm)
+ * - iterative orthogonal Procrustes updates on temp_mat
+ * - returns W = M_pca @ R
+ */
+ private static StiefelTransform runItqTrainer(
+ double[][] xHdNorm, // [N][D] normalized residuals
+ int quantizedDim, // d
+ Random rng,
+ int nTrainingIterations
+ ) {
+ logProgress("\t[stage] Starting runITQTrainer...");
+
+ final int N = xHdNorm.length;
+ final int D = xHdNorm[0].length;
+ final int d = quantizedDim;
+
+ // PCA basis M = V[:, :d] where X = U S V^T
+ RealMatrix X = new Array2DRowRealMatrix(xHdNorm, false); // [N×D]
+ logProgress("\t[stage] Starting Native SVD for PCA...");
+ // We only need V, so we tell LAPACK to skip U to save massive amounts of time/memory
+ RealMatrix V = computeNativeV(X);
+ logProgress("\t[stage] Completed Native SVD...");
+ RealMatrix M = V.getSubMatrix(0, D - 1, 0, d - 1); // [D×d]
+
+ // Early stopping based on stabilized binary codes (subject to max iterations)
+ RealMatrix prevXbin = null;
+ int stable = 0;
+ final double STOP_FRAC = 4e-3; // stop when <0.4% bits flip TODO setting
+ final int STOP_PATIENCE = 3; // for 3 consecutive epochs TODO setting
+
+ // X_norm = X_hd_norm @ M (N x d)
+ RealMatrix Xnorm = nativeMultiply(X, M); // [N×d]
+
+ // temp_mat ~ N(0,1) in R^{d×d}. Initial random rotation.
+ RealMatrix temp = gaussianMatrix(d, d, rng);
+
+ // Loop: epoch 0..iters inclusive; last epoch computes R but does not update temp_mat.
+ RealMatrix R = null;
+ logProgress("\t[stage] ITQ training iterations started...");
+ for (int epoch = 0; epoch <= nTrainingIterations; epoch++) {
+ R = orthogonalize(temp); // R = U @ V^T
+
+ if (epoch < nTrainingIterations) {
+ // Use native multiplication for Xtr = Xnorm @ R
+ RealMatrix Xtr = nativeMultiply(Xnorm, R); // [N×d]
+ RealMatrix Xbin = sign01(Xtr); // [N×d] in {+1,-1}
+
+ if (prevXbin != null) {
+ long changed = 0;
+ long total = (long) N * (long) d;
+ for (int i = 0; i < N; i++) {
+ for (int j = 0; j < d; j++) {
+ if (Xbin.getEntry(i, j) != prevXbin.getEntry(i, j)) changed++;
+ }
+ }
+ double frac = (double) changed / (double) total;
+
+ logProgress(String.format(java.util.Locale.ROOT,
+ "\titeration %d/%d - fracBitsChanged=%.6g (stable=%d/%d)",
+ epoch, nTrainingIterations, frac, stable, STOP_PATIENCE));
+
+ if (frac < STOP_FRAC) {
+ if (++stable >= STOP_PATIENCE) break;
+ } else {
+ stable = 0;
+ }
+ }
+
+ // Optimized temp = Xnorm.T @ Xbin
+ // This avoids the physical transpose of the large Xnorm matrix
+ prevXbin = Xbin;
+ temp = nativeMultiplyTransposedA(Xnorm, Xbin); // [d×d]
+ }
+ }
+ logProgress("\t[stage] ITQ training completed.");
+
+ // W = M @ R (shape [D×d])
+ // Use native multiply for the final transform matrix
+ RealMatrix W = nativeMultiply(M, R);
+ return new StiefelTransform(W);
+ }
+
+ /**
+ * LANDING trainer TODO incomplete
+ *
+ * NOTE: This is a DRAFT (not hardened) implementation using RealMatrix operations.
+ * We can SIMD/block-optimize the hot math once correctness is validated.
+ */
+ private static StiefelTransform runLandingTrainer(
+ double[][] xHdNorm, // [N][D] normalized residuals
+ int quantizedDim, // d
+ Random rng,
+ int nTrainingIterations
+ ) {
+ final int N = xHdNorm.length;
+ final int D = xHdNorm[0].length;
+ final int d = quantizedDim;
+
+ RealMatrix X = new Array2DRowRealMatrix(xHdNorm, false); // [N×D]
+
+ // PCA basis M_pca = V[:, :d]
+ SingularValueDecomposition svdX = new SingularValueDecomposition(X);
+ RealMatrix V = svdX.getV();
+ RealMatrix M_pca = V.getSubMatrix(0, D - 1, 0, d - 1); // [D×d]
+
+ // Init:
+ // temp(d×d) -> Q = U V^T
+ // A = Q @ M_pca^T (d×D)
+ // W = A^T (D×d)
+ RealMatrix Q = orthogonalize(gaussianMatrix(d, d, rng));
+ RealMatrix A = Q.multiply(M_pca.transpose()); // [d×D]
+ RealMatrix W = A.transpose(); // [D×d]
+
+ final RealMatrix I_D = MatrixUtils.createRealIdentityMatrix(D);
+ final double invSqrtD = 1.0 / Math.sqrt(d);
+
+ int[] idx = new int[N];
+ for (int i = 0; i < N; i++) idx[i] = i;
+
+ for (int epoch = 0; epoch < nTrainingIterations; epoch++) {
+ shuffleInPlace(idx, rng);
+
+ for (int off = 0; off < N; off += LANDING_BATCH_SIZE) {
+ final int end = Math.min(N, off + LANDING_BATCH_SIZE);
+ final int B = end - off;
+
+ // Build batch matrix X_batch (B×D)
+ double[][] xb = new double[B][D];
+ for (int bi = 0; bi < B; bi++) {
+ xb[bi] = xHdNorm[idx[off + bi]];
+ }
+ RealMatrix Xb = new Array2DRowRealMatrix(xb, false);
+
+ // X_sign = sign(Xb @ W) / sqrt(d)
+ RealMatrix Xproj = Xb.multiply(W); // [B×d]
+ RealMatrix Xsign = sign01(Xproj).scalarMultiply(invSqrtD);// [B×d]
+
+ // gradA = (X_sign^T X_sign) A - (X_sign^T X_batch)
+ RealMatrix XtX = Xsign.transpose().multiply(Xsign); // [d×d]
+ RealMatrix XtB = Xsign.transpose().multiply(Xb); // [d×D]
+ RealMatrix gradA = XtX.multiply(A).subtract(XtB); // [d×D]
+
+ // gradW = (X_batch^T X_batch) ( W (A A^T) - A^T )
+ RealMatrix G = Xb.transpose().multiply(Xb); // [D×D]
+ RealMatrix AAT = A.multiply(A.transpose()); // [d×d]
+ RealMatrix term = W.multiply(AAT).subtract(A.transpose());// [D×d]
+ RealMatrix gradW = G.multiply(term); // [D×d]
+
+ double lr = 0.1 * Math.pow(0.98, epoch);
+ double lambda = 1.0;
+
+ // skew = 0.5 * (A^T gradA - (A^T gradA)^T)
+ RealMatrix Mmat = A.transpose().multiply(gradA); // [D×D]
+ RealMatrix skew = Mmat.subtract(Mmat.transpose()).scalarMultiply(0.5);
+
+ // A -= lr * A @ ( lambda*(A^T A - I) + skew )
+ RealMatrix ortho = A.transpose().multiply(A).subtract(I_D);
+ RealMatrix updateCore = ortho.scalarMultiply(lambda).add(skew);
+ A = A.subtract(A.multiply(updateCore).scalarMultiply(lr));
+
+ // W -= lr * gradW
+ W = W.subtract(gradW.scalarMultiply(lr));
+ }
+ }
+
+ return new StiefelTransform(W);
+ }
+
+ private static void shuffleInPlace(int[] a, Random rng) {
+ for (int i = a.length - 1; i > 0; i--) {
+ int j = rng.nextInt(i + 1);
+ int tmp = a[i];
+ a[i] = a[j];
+ a[j] = tmp;
+ }
+ }
+
+ private static RealMatrix gaussianMatrix(int rows, int cols, Random rng) {
+ double[][] m = new double[rows][cols];
+ for (int i = 0; i < rows; i++) {
+ double[] r = m[i];
+ for (int j = 0; j < cols; j++) {
+ r[j] = rng.nextGaussian();
+ }
+ }
+ return new Array2DRowRealMatrix(m, false);
+ }
+
+ /** Orthonormal factor: U @ V^T from SVD(M). */
+// private static RealMatrix orthogonalize(RealMatrix m) {
+// SingularValueDecomposition svd = new SingularValueDecomposition(m);
+// return svd.getU().multiply(svd.getVT());
+// }
+
+ // Basic dgemm routine.
+ private static RealMatrix nativeMultiply(RealMatrix A, RealMatrix B) {
+ int m = A.getRowDimension();
+ int n = B.getColumnDimension();
+ int k = A.getColumnDimension(); // Must match B's row dimension
+
+ // 1. Convert to 1D Column-Major for BLAS
+ double[] aArr = serializeColumnMajor(A);
+ double[] bArr = serializeColumnMajor(B);
+ double[] cArr = new double[m * n];
+
+ // 2. Call Native DGEMM: C = alpha*A*B + beta*C
+ // "N", "N" means 'No Transpose' for both A and B
+ BLAS.getInstance().dgemm(
+ "N", "N",
+ m, n, k,
+ 1.0, // alpha
+ aArr, m, // A and its leading dimension
+ bArr, k, // B and its leading dimension
+ 0.0, // beta (ignores existing values in cArr)
+ cArr, m // C and its leading dimension
+ );
+
+ // 3. Convert back to RealMatrix
+ return new Array2DRowRealMatrix(deserializeColumnMajor(cArr, m, n), false);
+ }
+
+ // Uses the dgesvd routine but tells LAPACK to skip the U matrix (for PCA). Saves time....
+ private static RealMatrix computeNativeV(RealMatrix x) {
+ int rows = x.getRowDimension();
+ int cols = x.getColumnDimension();
+ double[] a = serializeColumnMajor(x);
+
+ double[] s = new double[Math.min(rows, cols)];
+ double[] vt = new double[cols * cols]; // Right singular vectors
+ intW info = new intW(0);
+
+ // Query optimal workspace size (standard LAPACK practice)
+ double[] workQuery = new double[1];
+ LAPACK.getInstance().dgesvd("N", "A", rows, cols, a, rows, s, null, 1, vt, cols, workQuery, -1, info);
+
+ int lwork = (int) workQuery[0];
+ double[] work = new double[lwork];
+
+ // Actual call: "N" = No U matrix, "A" = All columns of V^T
+ LAPACK.getInstance().dgesvd("N", "A", rows, cols, a, rows, s, null, 1, vt, cols, work, lwork, info);
+
+ if (info.val != 0) {
+ throw new RuntimeException("Native PCA SVD failed with info=" + info.val);
+ }
+
+ // Convert VT (column-major) back to RealMatrix
+ return new Array2DRowRealMatrix(deserializeColumnMajor(vt, cols, cols), false);
+ }
+
+ /** Computes A^T * B using Native BLAS dgemm */
+ private static RealMatrix nativeMultiplyTransposedA(RealMatrix A, RealMatrix B) {
+ int m = A.getColumnDimension(); // A.T has rows = A.cols
+ int n = B.getColumnDimension();
+ int k = A.getRowDimension(); // A.T has cols = A.rows
+
+ double[] aArr = serializeColumnMajor(A);
+ double[] bArr = serializeColumnMajor(B);
+ double[] cArr = new double[m * n];
+
+ // "T" = Transpose Matrix A, "N" = No Transpose Matrix B
+ BLAS.getInstance().dgemm(
+ "T", "N",
+ m, n, k,
+ 1.0,
+ aArr, k, // Leading dimension is k because A is currently [k x m]
+ bArr, k,
+ 0.0,
+ cArr, m
+ );
+
+ return new Array2DRowRealMatrix(deserializeColumnMajor(cArr, m, n), false);
+ }
+
+ private static double calculateMSE(RealMatrix a, RealMatrix b) {
+ double sumSq = 0;
+ int rows = a.getRowDimension();
+ int cols = a.getColumnDimension();
+ for (int i = 0; i < rows; i++) {
+ for (int j = 0; j < cols; j++) {
+ double diff = a.getEntry(i, j) - b.getEntry(i, j);
+ sumSq += diff * diff;
+ }
+ }
+ return sumSq / (rows * cols);
+ }
+
+ private static RealMatrix orthogonalize(RealMatrix m) {
+ int rows = m.getRowDimension();
+ int cols = m.getColumnDimension();
+
+ // Convert to column-major for Native LAPACK
+ double[] a = serializeColumnMajor(m);
+
+ // Prepare output buffers
+ double[] u = new double[rows * rows];
+ double[] vt = new double[cols * cols];
+ double[] s = new double[Math.min(rows, cols)];
+
+ // int wrapper required....
+ intW info = new intW(0);
+
+ // LAPACK Workspace (Required for dgesvd): Pass -1 for lwork and a 1-element array for work
+ double[] workQuery = new double[1];
+ LAPACK.getInstance().dgesvd(
+ "A", "A", rows, cols, a, rows, s, u, rows, vt, cols,
+ workQuery, -1, info
+ );
+
+ // Extract the optimal lwork and allocate the real workspace
+ int lwork = (int) workQuery[0];
+ double[] work = new double[lwork];
+
+ // Native SVD call
+ LAPACK.getInstance().dgesvd(
+ "A", "A", // Compute all columns of U and all rows of VT
+ rows, cols, // Matrix dimensions
+ a, rows, // Input matrix and its leading dimension
+ s, // Output singular values
+ u, rows, // Output U and its leading dimension
+ vt, cols, // Output VT and its leading dimension
+ work, lwork, // Workspace
+ info // Success/Fail code
+ );
+
+ if (info.val != 0) {
+ throw new RuntimeException("LAPACK dgesvd failed with info=" + info.val);
+ }
+
+ // Convert back to Commons Math
+ RealMatrix U = new Array2DRowRealMatrix(deserializeColumnMajor(u, rows, rows), false);
+ RealMatrix VT = new Array2DRowRealMatrix(deserializeColumnMajor(vt, cols, cols), false);
+
+ return U.multiply(VT);
+ }
+
+ /** Converts an Apache Commons RealMatrix to a column-major 1D array for LAPACK. */
+ private static double[] serializeColumnMajor(RealMatrix m) {
+ int rows = m.getRowDimension();
+ int cols = m.getColumnDimension();
+ double[] data = new double[rows * cols];
+ for (int c = 0; c < cols; c++) {
+ for (int r = 0; r < rows; r++) {
+ data[c * rows + r] = m.getEntry(r, c);
+ }
+ }
+ return data;
+ }
+
+ /** Converts a column-major 1D array back into a row-major double[][] for Commons Math. */
+ private static double[][] deserializeColumnMajor(double[] data, int rows, int cols) {
+ double[][] matrix = new double[rows][cols];
+ for (int c = 0; c < cols; c++) {
+ for (int r = 0; r < rows; r++) {
+ matrix[r][c] = data[c * rows + r];
+ }
+ }
+ return matrix;
+ }
+
+ /**
+ * Elementwise sign producing +/-1 (no zeros), consistent with ASH encoding convention:
+ * bit=1 iff projection > 0, else bit=0.
+ */
+ private static RealMatrix sign01(RealMatrix m) {
+ final int r = m.getRowDimension();
+ final int c = m.getColumnDimension();
+ double[][] out = new double[r][c];
+ for (int i = 0; i < r; i++) {
+ double[] row = out[i];
+ for (int j = 0; j < c; j++) {
+ row[j] = (m.getEntry(i, j) > 0.0) ? 1.0 : -1.0;
+ }
+ }
+ return new Array2DRowRealMatrix(out, false);
+ }
+
@Override
public int compressorSize() {
// Must exactly match write(out, version)
@@ -1379,6 +1903,11 @@ private static int landmarksHash(VectorFloat>[] a) {
return h;
}
+ private static void logProgress(String msg) {
+ System.out.println(msg);
+ System.out.flush();
+ }
+
@Override
public String toString() {
String optimizerName;
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 49baff0c7..a75cc68ab 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -44,6 +44,30 @@ private static void printScorerInfo(String label, ASHBlockScorer scorer) {
);
}
+ private static void logProgress(String msg) {
+ System.out.println(msg);
+ System.out.flush();
+ }
+
+ private static int parseOptimizerFromProperty() {
+ String opt = System.getProperty("jvector.ash.optimizer", "random").trim().toLowerCase();
+
+ if ("random".equals(opt)) return AsymmetricHashing.RANDOM;
+ if ("itq".equals(opt)) return AsymmetricHashing.ITQ;
+ if ("landing".equals(opt)) return AsymmetricHashing.LANDING;
+
+ throw new IllegalArgumentException(
+ "Unknown jvector.ash.optimizer=" + opt + " (expected random|itq|landing)"
+ );
+ }
+
+ private static String optimizerName(int optimizer) {
+ if (optimizer == AsymmetricHashing.RANDOM) return "RANDOM";
+ if (optimizer == AsymmetricHashing.ITQ) return "ITQ";
+ if (optimizer == AsymmetricHashing.LANDING) return "LANDING";
+ return "UNKNOWN(" + optimizer + ")";
+ }
+
public static void testASHEncodings(String filenameBase, String filenameQueries) throws IOException {
// ------------------------------------------------------------
// Benchmark configuration (runtime flags)
@@ -52,7 +76,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "true"));
final boolean RUN_SCALAR_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
@@ -149,15 +173,28 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
vectors.size(), queries.size(), dimension);
// ------------------------------------------------------------------
- // Build ASH
+ // Build ASH (centroids + optional training), then encode
// ------------------------------------------------------------------
var ravv = new ListRandomAccessVectorValues(vectors, dimension);
- var ash = AsymmetricHashing.initialize(ravv, AsymmetricHashing.RANDOM, encodedBits, landmarkCount);
+ // Choose optimizer at runtime:
+ // -Djvector.ash.optimizer=random|itq|landing
+ final int optimizer = parseOptimizerFromProperty();
+
+ logProgress("\tASH optimizer = " + optimizerName(optimizer));
+
+ logProgress("\t[stage] ASH initialize: starting (centroids + training)");
+ long initStart = System.nanoTime();
+ var ash = AsymmetricHashing.initialize(ravv, optimizer, encodedBits, landmarkCount);
+ long initEnd = System.nanoTime();
+ logProgress("\t[stage] ASH initialize: done in " + (initEnd - initStart) / 1e9 + " seconds");
+
+ logProgress("\t[stage] ASH encodeAll: starting");
long startTime = System.nanoTime();
CompressedVectors ashVecs = ash.encodeAll(ravv);
long endTime = System.nanoTime();
- System.out.println("\tEncoding took " + (endTime - startTime) / 1e9 + " seconds");
+ logProgress("\t[stage] ASH encodeAll: done in " + (endTime - startTime) / 1e9 + " seconds");
+
double encSeconds = (endTime - startTime) / 1e9;
double encThroughput = finalVectorCount / encSeconds;
diff --git a/pom.xml b/pom.xml
index 7cca5b378..034459823 100644
--- a/pom.xml
+++ b/pom.xml
@@ -1,4 +1,4 @@
-
4.0.0
@@ -42,6 +42,7 @@
UTF-8
true
4.0.0-rc.7-SNAPSHOT
+ 3.0.3
jvector-base
@@ -198,6 +199,31 @@
+
+ com.github.fommil.netlib
+ core
+ 1.1.2
+
+
+ dev.ludovic.netlib
+ netlib
+ ${ludovic.netlib.version}
+
+
+ dev.ludovic.netlib
+ blas
+ ${ludovic.netlib.version}
+
+
+ dev.ludovic.netlib
+ lapack
+ ${ludovic.netlib.version}
+
+
+ dev.ludovic.netlib
+ arpack
+ ${ludovic.netlib.version}
+
org.junit.jupiter
junit-jupiter-engine
From 6b738859abaff1bb407eb38c867e671148619e98 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Thu, 8 Jan 2026 20:27:19 +0000
Subject: [PATCH 026/243] Added k-recall@R computations to DistancesASH and
added a truncated ivecs loader to SiftLoader.
---
.../quantization/AsymmetricHashing.java | 2 +-
.../jbellis/jvector/example/DistancesASH.java | 133 +++++++++++++++---
.../jvector/example/util/SiftLoader.java | 42 ++++++
3 files changed, 153 insertions(+), 24 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index ddd02da3c..652f42aad 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -72,7 +72,7 @@ public class AsymmetricHashing implements VectorCompressor> vectors = SiftLoader.readFvecs(filenameBase, nVectors);
List> queries = SiftLoader.readFvecs(filenameQueries, nQueries);
+ List groundTruth = SiftLoader.readIvecsAsArrays(filenameGT, nQueries);
// ------------------------------------------------------------------
// Remove zero-norm vectors (undefined for ASH / Eq. 11)
// ------------------------------------------------------------------
+
int beforeBase = vectors.size();
- List> filteredBase = new ArrayList<>(vectors.size());
- for (VectorFloat> v : vectors) {
+
+ // Maps original index to filtered index. Value is -1 if vector was removed.
+ int[] oldToNew = new int[beforeBase];
+ java.util.Arrays.fill(oldToNew, -1);
+
+ List> filteredBase = new ArrayList<>(beforeBase);
+ for (int i = 0; i < beforeBase; i++) {
+ VectorFloat> v = vectors.get(i);
if (VectorUtil.dotProduct(v, v) > 0f) {
+ oldToNew[i] = filteredBase.size();
filteredBase.add(v);
}
}
@@ -149,7 +164,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = (dimension / 4) + HEADER_BITS;
+ int encodedBits = (dimension / 4) + HEADER_BITS - 64; // (dimension / 4) + HEADER_BITS;
// Payload must be 64-bit aligned for SIMD
int payloadBits = encodedBits - HEADER_BITS;
if ((payloadBits & 63) != 0) {
@@ -337,6 +352,69 @@ public static void testASHEncodings(String filenameBase, String filenameQueries)
System.out.println("\tAverage absolute dot-product error = " + distanceError);
}
+ // ==================================================================
+ // [1b] Recall@K run (Parallelized)
+ // ==================================================================
+ if (RUN_RECALL_CHECK) {
+ int[] atValues = {10, 15, 20, 30, 40, 50};
+ int maxAt = 50;
+ List> recallTasks = new ArrayList<>();
+
+ logProgress("\t[stage] Computing " + RECALL_K + "-Recall@K...");
+ for (int start = 0; start < finalQueryCount; start += chunkSize) {
+ final int s = start;
+ final int e = Math.min(start + chunkSize, finalQueryCount);
+
+ recallTasks.add(simdExecutor.submit(() -> {
+ double[] localTotalRecall = new double[atValues.length];
+ for (int i = s; i < e; i++) {
+ VectorFloat> q = finalQueries.get(i);
+ ScoreFunction.ApproximateScoreFunction f = ashVecsFinal.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+
+ // Min-heap to keep top 50 results (storing score as int bits and ordinal as long)
+ var topCandidates = new java.util.PriorityQueue((a, b) -> Float.compare(Float.intBitsToFloat((int) a[0]), Float.intBitsToFloat((int) b[0])));
+ for (int j = 0; j < finalVectorCount; j++) {
+ float score = f.similarityTo(j);
+ if (topCandidates.size() < maxAt) {
+ topCandidates.add(new long[]{Float.floatToRawIntBits(score), j});
+ } else if (score > Float.intBitsToFloat((int) topCandidates.peek()[0])) {
+ topCandidates.poll();
+ topCandidates.add(new long[]{Float.floatToRawIntBits(score), j});
+ }
+ }
+
+ int[] topIndices = new int[topCandidates.size()];
+ for (int rank = topCandidates.size() - 1; rank >= 0; rank--) topIndices[rank] = (int) topCandidates.poll()[1];
+
+ int[] queryGT = groundTruth.get(i);
+ for (int aIdx = 0; aIdx < atValues.length; aIdx++) {
+ int at = atValues[aIdx];
+ java.util.Set topAtSet = new java.util.HashSet<>();
+ for (int r = 0; r < Math.min(at, topIndices.length); r++) {
+ topAtSet.add((newToOldFinal == null) ? topIndices[r] : newToOldFinal[topIndices[r]]);
+ }
+ int matches = 0;
+ for (int g = 0; g < RECALL_K; g++) {
+ if (topAtSet.contains(queryGT[g])) matches++;
+ }
+ localTotalRecall[aIdx] += (double) matches / RECALL_K;
+ }
+ }
+ return localTotalRecall;
+ }));
+ }
+
+ double[] totalRecall = new double[atValues.length];
+ for (ForkJoinTask t : recallTasks) {
+ double[] local = t.join();
+ for (int aIdx = 0; aIdx < atValues.length; aIdx++) totalRecall[aIdx] += local[aIdx];
+ }
+
+ for (int aIdx = 0; aIdx < atValues.length; aIdx++) {
+ System.out.format("\tASH %d-recall@%d = %.4f%n", RECALL_K, atValues[aIdx], totalRecall[aIdx] / finalQueryCount);
+ }
+ }
+
// ==================================================================
// [2] Scalar ASH scoring timing (reference)
// ==================================================================
@@ -522,7 +600,8 @@ public static void runSIFT() throws IOException {
var baseVectors = "siftsmall/siftsmall_base.fvecs";
var queryVectors = "siftsmall/siftsmall_query.fvecs";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runGIST() throws IOException {
@@ -530,7 +609,8 @@ public static void runGIST() throws IOException {
var baseVectors = "./fvec/gist/gist_base.fvecs";
var queryVectors = "./fvec/gist/gist_query.fvecs";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runCohere100k() throws IOException {
@@ -538,7 +618,8 @@ public static void runCohere100k() throws IOException {
var baseVectors = "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_base_vectors_100000.fvec";
var queryVectors = "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_query_vectors_10000.fvec";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "./fvec/cohere-100k/cohere_embed-english-v3.0_1024_indices_b100000_q10000_k100.ivec";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runADA() throws IOException {
@@ -546,7 +627,8 @@ public static void runADA() throws IOException {
var baseVectors = "./fvec/ada-002/ada_002_100000_base_vectors.fvec";
var queryVectors = "./fvec/ada-002/ada_002_100000_query_vectors_10000.fvec";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "./fvec/ada-002/ada_002_100000_indices_query_10000.ivec";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runColbert() throws IOException {
@@ -554,7 +636,8 @@ public static void runColbert() throws IOException {
var baseVectors = "./fvec/wikipedia_squad/1M/colbertv2.0_128_base_vectors_1000000.fvec";
var queryVectors = "./fvec/wikipedia_squad/1M/colbertv2.0_128_query_vectors_100000.fvec";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runOpenai1536() throws IOException {
@@ -562,7 +645,8 @@ public static void runOpenai1536() throws IOException {
var baseVectors = "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_base_vectors.fvec";
var queryVectors = "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_query_vectors_10000.fvec";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "./fvec/openai-v3-large-1536-100k/text-embedding-3-large_1536_100000_indices_query_10000.ivec";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runOpenai3072() throws IOException {
@@ -570,7 +654,8 @@ public static void runOpenai3072() throws IOException {
var baseVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_base_vectors.fvec";
var queryVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_query_vectors_10000.fvec";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "./fvec/openai-v3-large-3072-100k/text-embedding-3-large_3072_100000_indices_query_10000.ivec";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runCap6m() throws IOException {
@@ -578,7 +663,8 @@ public static void runCap6m() throws IOException {
var baseVectors = "./fvec/cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_base_6m_norm_shuffle.fvecs";
var queryVectors = "./fvec/cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_query_10k_norm_shuffle.fvecs";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "./fvec/cap-6m/cap_6m_gt_norm_shuffle_ip_k100.ivecs";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void runCohere10m() throws IOException {
@@ -586,18 +672,19 @@ public static void runCohere10m() throws IOException {
var baseVectors = "./fvec/cohere-10m/cohere_wiki_en_flat_base_10m_norm.fvecs";
var queryVectors = "./fvec/cohere-10m/cohere_wiki_en_flat_query_10k_norm.fvecs";
- testASHEncodings(baseVectors, queryVectors);
+ var gtVectors = "./fvec/cohere-10m/cohere_wiki_en_flat_gt_10m_ip_k100.ivecs";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
}
public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
- runCohere100k();
- runADA();
- runOpenai1536();
- runOpenai3072();
-// runCap6m();
-// runCohere10m();
+// runCohere100k();
+// runADA();
+// runOpenai1536();
+// runOpenai3072();
+ runCap6m();
+ runCohere10m();
}
}
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 cccfc1cea..f68ae340d 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
@@ -145,6 +145,48 @@ public static List> readFvecs(String filePath, int numVectors)
return vectors;
}
+ /**
+ * Read all ground-truth neighbors from an .ivecs file as a list of primitive arrays.
+ */
+ public static List readIvecsAsArrays(String filename) {
+ try {
+ return readIvecsAsArrays(filename, Integer.MAX_VALUE);
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+ }
+
+ /**
+ * Read at most {@code numVectors} from an .ivecs file as a list of primitive arrays.
+ */
+ public static List readIvecsAsArrays(String filePath, int numVectors) throws IOException {
+ if (numVectors <= 0) {
+ throw new IllegalArgumentException("numVectors must be > 0, got " + numVectors);
+ }
+
+ var groundTruth = new ArrayList(Math.min(numVectors, 1024));
+
+ try (var dis = new DataInputStream(new BufferedInputStream(new FileInputStream(filePath)))) {
+ int count = 0;
+ while (count < numVectors && dis.available() > 0) {
+ int k = Integer.reverseBytes(dis.readInt());
+ int[] neighbors = new int[k];
+ for (int i = 0; i < k; i++) {
+ neighbors[i] = Integer.reverseBytes(dis.readInt());
+ }
+ groundTruth.add(neighbors);
+ count++;
+ }
+ }
+
+ if (numVectors != Integer.MAX_VALUE && groundTruth.size() != numVectors) {
+ throw new IOException(String.format("File %s ended early: requested %d, read %d",
+ filePath, numVectors, groundTruth.size()));
+ }
+
+ return groundTruth;
+ }
+
/**
* Read all ground-truth neighbors from an .ivecs file.
*/
From 1582dbff94e937c088f0158d26a4475088e5adde Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Thu, 8 Jan 2026 22:01:13 +0000
Subject: [PATCH 027/243] Reduce allocations in quantizeBody by using a
workspace for several variables.
---
.../quantization/AsymmetricHashing.java | 33 ++++++++++++++++---
1 file changed, 28 insertions(+), 5 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 652f42aad..6e3b1cb48 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -621,6 +621,23 @@ void quantizeBody(VectorFloat> vector,
}
public static final class RandomBinaryQuantizer implements BinaryQuantizer {
+ // ThreadLocal to hold the reusable buffers.
+ // Using a simple container class to avoid multiple ThreadLocal lookups.
+ private static final ThreadLocal WORKSPACE = ThreadLocal.withInitial(Workspace::new);
+
+ private static final class Workspace {
+ float[] x;
+ float[] muArr;
+ float[] xhat;
+
+ void ensureCapacity(int dim) {
+ if (x == null || x.length < dim) {
+ x = new float[dim];
+ muArr = new float[dim];
+ xhat = new float[dim];
+ }
+ }
+ }
@Override
public void quantizeBody(VectorFloat> vector,
@@ -633,14 +650,21 @@ public void quantizeBody(VectorFloat> vector,
final float[][] A = stiefel.AFloat;
final int originalDim = stiefel.cols;
+ // Get workspace for this thread and ensure correct sizing (originalDim may have changed)
+ Workspace ws = WORKSPACE.get();
+ ws.ensureCapacity(originalDim);
+
+ // Alias for readability (no new allocations)
+ final float[] x = ws.x;
+ final float[] muArr = ws.muArr;
+ final float[] xhat = ws.xhat;
+
final var vecUtil =
io.github.jbellis.jvector.vector.VectorizationProvider
.getInstance()
.getVectorUtilSupport();
- // Copy vector and mean once (no virtual calls in inner loop)
- final float[] x = new float[originalDim];
- final float[] muArr = new float[originalDim];
+ // Copy vector and mean once into workspace
for (int d = 0; d < originalDim; d++) {
x[d] = vector.get(d);
muArr[d] = mu.get(d);
@@ -649,8 +673,7 @@ public void quantizeBody(VectorFloat> vector,
// ASH paper, Eq. 6, normalization
final float invNorm = (residualNorm > 0f) ? (1.0f / residualNorm) : 0.0f;
- // Normalized residual x̂
- final float[] xhat = new float[originalDim];
+ // Compute normalized residual x̂
for (int d = 0; d < originalDim; d++) {
xhat[d] = (x[d] - muArr[d]) * invNorm;
}
From d8097e1f0874b85419c32c6c96fef4182f622419 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Mon, 12 Jan 2026 01:18:37 +0000
Subject: [PATCH 028/243] Replaced QR with Thin SVD in runWithinTraining and
temporarily overriding payload alignment exception for recall benchmarking
purposes.
---
.../quantization/AsymmetricHashing.java | 175 ++++++++++--------
.../jbellis/jvector/example/DistancesASH.java | 53 ++++--
.../jbellis/jvector/example/DistancesPQ.java | 12 +-
3 files changed, 146 insertions(+), 94 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 6e3b1cb48..f5444e608 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -372,15 +372,15 @@ private static void validateEncodedBits(
);
}
- if ((payloadBits & 63) != 0) {
- throw new IllegalArgumentException(
- "Invalid encodedBits=" + encodedBits +
- ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
- "for aligned binary payload. " +
- "Got payloadBits=" + payloadBits +
- " (headerBits=" + headerBits + ")."
- );
- }
+// if ((payloadBits & 63) != 0) {
+// throw new IllegalArgumentException(
+// "Invalid encodedBits=" + encodedBits +
+// ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
+// "for aligned binary payload. " +
+// "Got payloadBits=" + payloadBits +
+// " (headerBits=" + headerBits + ")."
+// );
+// }
if (payloadBits > originalDim) {
throw new IllegalArgumentException(
@@ -999,56 +999,78 @@ public StiefelTransform(RealMatrix W) {
/**
* RANDOM (untrained) Stiefel initialization.
- * Generates a random orthonormal projection from R^originalDim → R^quantizedDim
- * by sampling a Gaussian matrix and orthonormalizing via QR decomposition.
- * A deterministic sign normalization is applied so results are reproducible
- * across runs and linear algebra backends.
- * @param originalDim dimensionality of the input vectors
- * @param quantizedDim dimensionality of the binary embedding
- * @param rng random number generator
+ *
+ * Samples a Gaussian matrix G ∈ R^{D×d} and computes its polar factor
+ * W = U·Vᵀ via a thin SVD (G = U·S·Vᵀ). The resulting W has orthonormal
+ * columns (WᵀW = I) and can be used as an orthonormal projection matrix
+ * (e.g., y = Wᵀx maps R^D → R^d).
+ *
+ *
Applies a deterministic per-column sign alignment to reduce SVD sign
+ * ambiguity by flipping columns of W to be positively correlated with the
+ * corresponding columns of G.
+ *
+ * @param D dimensionality of the input vectors
+ * @param d dimensionality of the projected (embedding) space
+ * @param rng random number generator used to sample G
+ * @return a {@link StiefelTransform} containing W ∈ R^{D×d} with orthonormal columns
+ * @throws IllegalArgumentException if d <= 0 or d > D
+ * @throws RuntimeException if the native SVD fails
*/
- public static StiefelTransform runWithoutTraining(int originalDim,
- int quantizedDim,
- Random rng) {
-
- if (quantizedDim <= 0 || quantizedDim > originalDim) {
- throw new IllegalArgumentException(
- "Invalid quantizedDim=" + quantizedDim +
- " for originalDim=" + originalDim
- );
+ public static StiefelTransform runWithoutTraining(int D, int d, Random rng) {
+ if (d <= 0 || d > D) {
+ throw new IllegalArgumentException("Invalid d=" + d + " for D=" + D);
}
- // Sample Gaussian matrix G ∈ R^{originalDim × quantizedDim}
- double[][] gaussian = new double[originalDim][quantizedDim];
- for (int i = 0; i < originalDim; i++) {
- for (int j = 0; j < quantizedDim; j++) {
- gaussian[i][j] = rng.nextGaussian();
- }
+ // Sample Gaussian matrix G [D x d]
+ // Using a 1D array for native speed, then wrapping it
+ double[] gData = new double[D * d];
+ for (int i = 0; i < gData.length; i++) {
+ gData[i] = rng.nextGaussian();
}
+ RealMatrix G = new Array2DRowRealMatrix(deserializeColumnMajor(gData, D, d), false);
+
+ // Perform Thin SVD: G = U S VT
+ // We need U [D x d] and VT [d x d]
+ double[] u = new double[D * d];
+ double[] vt = new double[d * d];
+ double[] s = new double[d];
+ intW info = new intW(0);
+
+ // Workspace query for dgesvd
+ double[] workQuery = new double[1];
+ LAPACK.getInstance().dgesvd("S", "S", D, d, gData, D, s, u, D, vt, d, workQuery, -1, info);
+ double[] work = new double[(int) workQuery[0]];
- RealMatrix G = MatrixUtils.createRealMatrix(gaussian);
+ // Execute dgesvd with Job "S" (Thin SVD)
+ LAPACK.getInstance().dgesvd("S", "S", D, d, gData, D, s, u, D, vt, d, work, work.length, info);
- // QR decomposition: G = Q R
- QRDecomposition qr = new QRDecomposition(G);
+ if (info.val != 0) {
+ throw new RuntimeException("Native SVD failed with info=" + info.val);
+ }
- // Commons Math returns full Q (m×m). We need thin Q (m×n).
- RealMatrix Qfull = qr.getQ(); // originalDim × originalDim
- RealMatrix Q = Qfull.getSubMatrix(0, originalDim - 1, 0, quantizedDim - 1); // originalDim × quantizedDim
+ // Compute W = U @ VT [D x d]
+ // This is the Polar Decomposition (nearest orthonormal matrix to G)
+ double[] wData = new double[D * d];
+ BLAS.getInstance().dgemm("N", "N", D, d, d, 1.0, u, D, vt, d, 0.0, wData, D);
- RealMatrix R = qr.getR(); // originalDim × quantizedDim
+ RealMatrix W = new Array2DRowRealMatrix(deserializeColumnMajor(wData, D, d), false);
- // Normalize QR sign ambiguity for determinism (RANDOM only)
- int diag = Math.min(R.getRowDimension(), R.getColumnDimension()); // == quantizedDim
- for (int j = 0; j < diag; j++) {
- if (R.getEntry(j, j) < 0.0) {
- for (int i = 0; i < Q.getRowDimension(); i++) {
- Q.setEntry(i, j, -Q.getEntry(i, j));
+ // Sign Alignment (Crucial for ASH/Random bits)
+ // To ensure the "direction" of our bits matches the Gaussian seeds,
+ // we flip columns of W so they are positively correlated with G.
+ for (int j = 0; j < d; j++) {
+ double dot = 0;
+ for (int i = 0; i < D; i++) {
+ dot += W.getEntry(i, j) * G.getEntry(i, j);
+ }
+ if (dot < 0.0) {
+ for (int i = 0; i < D; i++) {
+ W.setEntry(i, j, -W.getEntry(i, j));
}
}
}
- // W := Q (originalDim × quantizedDim)
- return new StiefelTransform(Q);
+ return new StiefelTransform(W);
}
@Override
@@ -1607,52 +1629,52 @@ private static double calculateMSE(RealMatrix a, RealMatrix b) {
}
private static RealMatrix orthogonalize(RealMatrix m) {
- int rows = m.getRowDimension();
- int cols = m.getColumnDimension();
+ int rows = m.getRowDimension(); // D
+ int cols = m.getColumnDimension(); // d
+ int k = Math.min(rows, cols); // The rank for thin SVD
- // Convert to column-major for Native LAPACK
+ // Move to column-major 1D array
double[] a = serializeColumnMajor(m);
- // Prepare output buffers
- double[] u = new double[rows * rows];
- double[] vt = new double[cols * cols];
- double[] s = new double[Math.min(rows, cols)];
-
- // int wrapper required....
+ // Prepare Thin-SVD buffers
+ // u: [rows * k] instead of [rows * rows]
+ // vt: [k * cols] instead of [cols * cols]
+ double[] u = new double[rows * k];
+ double[] vt = new double[k * cols];
+ double[] s = new double[k];
intW info = new intW(0);
- // LAPACK Workspace (Required for dgesvd): Pass -1 for lwork and a 1-element array for work
+ // Workspace Query
double[] workQuery = new double[1];
LAPACK.getInstance().dgesvd(
- "A", "A", rows, cols, a, rows, s, u, rows, vt, cols,
+ "S", "S", rows, cols, a, rows, s, u, rows, vt, k,
workQuery, -1, info
);
- // Extract the optimal lwork and allocate the real workspace
int lwork = (int) workQuery[0];
double[] work = new double[lwork];
- // Native SVD call
+ // Actual Thin-SVD call (Job "S")
LAPACK.getInstance().dgesvd(
- "A", "A", // Compute all columns of U and all rows of VT
- rows, cols, // Matrix dimensions
- a, rows, // Input matrix and its leading dimension
- s, // Output singular values
- u, rows, // Output U and its leading dimension
- vt, cols, // Output VT and its leading dimension
- work, lwork, // Workspace
- info // Success/Fail code
+ "S", "S", rows, cols, a, rows, s, u, rows, vt, k,
+ work, lwork, info
);
if (info.val != 0) {
throw new RuntimeException("LAPACK dgesvd failed with info=" + info.val);
}
- // Convert back to Commons Math
- RealMatrix U = new Array2DRowRealMatrix(deserializeColumnMajor(u, rows, rows), false);
- RealMatrix VT = new Array2DRowRealMatrix(deserializeColumnMajor(vt, cols, cols), false);
+ // Native Multiplication: Result = U @ VT
+ // Using dgemm here)
+ double[] resultData = new double[rows * cols];
+ BLAS.getInstance().dgemm(
+ "N", "N", rows, cols, k,
+ 1.0, u, rows, // A = u, lda = rows
+ vt, k, // B = vt, ldb = k
+ 0.0, resultData, rows // C = resultData, ldc = rows
+ );
- return U.multiply(VT);
+ return new Array2DRowRealMatrix(deserializeColumnMajor(resultData, rows, cols), false);
}
/** Converts an Apache Commons RealMatrix to a column-major 1D array for LAPACK. */
@@ -1660,9 +1682,14 @@ private static double[] serializeColumnMajor(RealMatrix m) {
int rows = m.getRowDimension();
int cols = m.getColumnDimension();
double[] data = new double[rows * cols];
+
+ // Accept RealMatrix double[][] implementation
+ double[][] raw = m.getData();
+
for (int c = 0; c < cols; c++) {
+ int colOffset = c * rows;
for (int r = 0; r < rows; r++) {
- data[c * rows + r] = m.getEntry(r, c);
+ data[colOffset + r] = raw[r][c];
}
}
return data;
@@ -1672,13 +1699,13 @@ private static double[] serializeColumnMajor(RealMatrix m) {
private static double[][] deserializeColumnMajor(double[] data, int rows, int cols) {
double[][] matrix = new double[rows][cols];
for (int c = 0; c < cols; c++) {
+ int colOffset = c * rows;
for (int r = 0; r < rows; r++) {
- matrix[r][c] = data[c * rows + r];
+ matrix[r][c] = data[colOffset + r];
}
}
return matrix;
}
-
/**
* Elementwise sign producing +/-1 (no zeros), consistent with ASH encoding convention:
* bit=1 iff projection > 0, else bit=0.
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 0862a6f58..f4f9a34ba 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -76,7 +76,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_RECALL_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.recall", "false"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.recall", "true"));
final int RECALL_K =
Integer.getInteger("jvector.bench.recall.k", 10);
@@ -100,8 +100,8 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
final int landmarkCount = 1;
// Define the benchmark size
- int maxQueries = 10_000;
- int maxVectors = 10_000_000;
+ int maxQueries = 1_000;
+ int maxVectors = 1_000_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -164,17 +164,17 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = (dimension / 4) + HEADER_BITS - 64; // (dimension / 4) + HEADER_BITS;
+ int encodedBits = dimension / 4; // 384; // (dimension / 4) + HEADER_BITS; // (dimension / 4) + HEADER_BITS;
// Payload must be 64-bit aligned for SIMD
int payloadBits = encodedBits - HEADER_BITS;
- if ((payloadBits & 63) != 0) {
- throw new IllegalArgumentException(
- "ASH payloadBits must be 64-bit aligned for SIMD. " +
- "Got payloadBits=" + payloadBits +
- " (encodedBits=" + encodedBits +
- ", HEADER_BITS=" + HEADER_BITS + ")"
- );
- }
+// if ((payloadBits & 63) != 0) {
+// throw new IllegalArgumentException(
+// "ASH payloadBits must be 64-bit aligned for SIMD. " +
+// "Got payloadBits=" + payloadBits +
+// " (encodedBits=" + encodedBits +
+// ", HEADER_BITS=" + HEADER_BITS + ")"
+// );
+// }
System.out.println(
"\toriginalDim = " + dimension +
@@ -394,8 +394,13 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
topAtSet.add((newToOldFinal == null) ? topIndices[r] : newToOldFinal[topIndices[r]]);
}
int matches = 0;
+ java.util.HashSet gtSeen = new java.util.HashSet<>(RECALL_K * 2);
+ // Compute and don't count duplicates (if present in GT or retrieved IDs)
for (int g = 0; g < RECALL_K; g++) {
- if (topAtSet.contains(queryGT[g])) matches++;
+ int gtId = queryGT[g];
+ if (gtSeen.add(gtId) && topAtSet.contains(gtId)) {
+ matches++;
+ }
}
localTotalRecall[aIdx] += (double) matches / RECALL_K;
}
@@ -623,7 +628,7 @@ public static void runCohere100k() throws IOException {
}
public static void runADA() throws IOException {
- System.out.println("Running ada_002");
+ System.out.println("Running ada-002");
var baseVectors = "./fvec/ada-002/ada_002_100000_base_vectors.fvec";
var queryVectors = "./fvec/ada-002/ada_002_100000_query_vectors_10000.fvec";
@@ -631,6 +636,15 @@ public static void runADA() throws IOException {
testASHEncodings(baseVectors, queryVectors, gtVectors);
}
+ public static void runADANoZeros() throws IOException {
+ System.out.println("Running ada-002-no-zeros");
+
+ var baseVectors = "./fvec/ada-002-no-zeros/ada_002_100000_base_vectors_no_zeros.fvec";
+ var queryVectors = "./fvec/ada-002-no-zeros/ada_002_100000_query_vectors_10000_no_zeros.fvec";
+ var gtVectors = "./fvec/ada-002-no-zeros/ada-002_gt_no_zeros.ivec";
+ testASHEncodings(baseVectors, queryVectors, gtVectors);
+ }
+
public static void runColbert() throws IOException {
System.out.println("Running colbertv2");
@@ -680,11 +694,12 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
-// runCohere100k();
+ runCohere100k();
// runADA();
-// runOpenai1536();
-// runOpenai3072();
- runCap6m();
- runCohere10m();
+ runADANoZeros();
+ runOpenai1536();
+ runOpenai3072();
+// runCap6m();
+// runCohere10m();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
index 47c88f975..df415d326 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -336,6 +336,15 @@ public static void runADA() throws IOException {
);
}
+ public static void runADANoZeros() throws IOException {
+ System.out.println("Running ada-002-no-zeros");
+
+ testPQEncodings(
+ "./fvec/ada-002-no-zeros/ada_002_100000_base_vectors_no_zeros.fvec",
+ "./fvec/ada-002-no-zeros/ada_002_100000_query_vectors_10000_no_zeros.fvec"
+ );
+ }
+
public static void runOpenai1536() throws IOException {
System.out.println("Running text-embedding-3-large_1536");
testPQEncodings(
@@ -370,7 +379,8 @@ public static void runCohere10m() throws IOException {
public static void main(String[] args) throws IOException {
runCohere100k();
- runADA();
+// runADA();
+ runADANoZeros();
runOpenai1536();
runOpenai3072();
// runCap6m();
From 2fad117b767fa0570d5b7b003e2e81342606e770 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Tue, 13 Jan 2026 05:45:49 +0000
Subject: [PATCH 029/243] Changed computeNativeV(X) to computeNativeVt(X) to
fix ITQ. Other small fixes. ITQ and multi-landmarks work.
---
.../quantization/AsymmetricHashing.java | 141 ++++++++++--------
.../jbellis/jvector/example/DistancesASH.java | 18 +--
2 files changed, 88 insertions(+), 71 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index f5444e608..79e0d327e 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -43,6 +43,16 @@
import org.apache.commons.math3.linear.SingularValueDecomposition;
import org.netlib.util.intW;
+import java.io.BufferedInputStream;
+import java.io.DataInputStream;
+import java.io.FileInputStream;
+import java.io.IOException;
+
+import java.io.BufferedOutputStream;
+import java.io.FileOutputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED;
/**
@@ -246,7 +256,7 @@ private AsymmetricHashing(int originalDim,
public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
int optimizer,
int encodedBits,
- int landmarkCount) {
+ int landmarkCount) throws IOException {
final int quantizedDim = encodedBits - HEADER_BITS;
@@ -265,20 +275,37 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
// payload bits are always 64-bit aligned by validateEncodedBits()
validateEncodedBits(encodedBits, HEADER_BITS, originalDim);
- // NOTE: points are treated as read-only by KMeansPlusPlusClusterer. Materialize once.
+ // NOTE: points are treated as read-only by KMeansPlusPlusClusterer.
+ // Materialize and L2-normalize points
VectorFloat>[] points = new VectorFloat>[ravvCopy.size()];
for (int i = 0; i < ravvCopy.size(); i++) {
- points[i] = ravvCopy.getVector(i);
+ VectorFloat> v = ravvCopy.getVector(i).copy(); // copy to avoid mutating original source
+ float norm = (float) Math.sqrt(VectorUtil.dotProduct(v, v));
+ if (norm > 1e-10f) {
+ VectorUtil.scale(v, 1.0f / norm);
+ }
+ points[i] = v;
}
VectorFloat>[] landmarks;
if (landmarkCount == 1) {
- // Mean centroid
- VectorFloat> mu0 = vectorTypeSupport.createFloatVector(originalDim);
+ double[] muAccumulator = new double[originalDim];
+
+ // Accumulate in double precision
for (int i = 0; i < points.length; i++) {
- VectorUtil.addInPlace(mu0, points[i]);
+ VectorFloat> p = points[i];
+ for (int d = 0; d < originalDim; d++) {
+ muAccumulator[d] += p.get(d);
+ }
}
- VectorUtil.scale(mu0, 1.0f / points.length);
+
+ // Scale and materialize back into a float vector
+ VectorFloat> mu0 = vectorTypeSupport.createFloatVector(originalDim);
+ double invN = 1.0 / points.length;
+ for (int d = 0; d < originalDim; d++) {
+ mu0.set(d, (float) (muAccumulator[d] * invN));
+ }
+
landmarks = new VectorFloat>[] { mu0 };
} else {
long kmStart = System.nanoTime();
@@ -305,8 +332,8 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
}
}
- // Keep determinism consistent within ASH. TODO define seed constant or remove.
- final Random rng = new Random(42);
+ // Keep determinism consistent within ASH. TODO explore more robust random initialization
+ final Random rng = new Random(103);
StiefelTransform stiefelTransform;
if (optimizer == RANDOM) {
@@ -372,15 +399,15 @@ private static void validateEncodedBits(
);
}
-// if ((payloadBits & 63) != 0) {
-// throw new IllegalArgumentException(
-// "Invalid encodedBits=" + encodedBits +
-// ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
-// "for aligned binary payload. " +
-// "Got payloadBits=" + payloadBits +
-// " (headerBits=" + headerBits + ")."
-// );
-// }
+ if ((payloadBits & 63) != 0) {
+ throw new IllegalArgumentException(
+ "Invalid encodedBits=" + encodedBits +
+ ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
+ "for aligned binary payload. " +
+ "Got payloadBits=" + payloadBits +
+ " (headerBits=" + headerBits + ")."
+ );
+ }
if (payloadBits > originalDim) {
throw new IllegalArgumentException(
@@ -406,36 +433,39 @@ private static LandmarkProjections computeLandmarkProjections(
int originalDim
) {
final int C = landmarks.length;
- final float[][] A = stiefelTransform.AFloat; // [d][D]
-
- final var vecUtil =
- io.github.jbellis.jvector.vector.VectorizationProvider
- .getInstance()
- .getVectorUtilSupport();
+ // Use a double-precision matrix for the projection
+ final double[][] A = stiefelTransform.AData; // [d][D]
final float[][] landmarkProj = new float[C][quantizedDim];
final float[] landmarkProjSum = new float[C];
- final float[] muArr = new float[originalDim];
+ // Materialize μ as double[] to maintain precision during the dot product
+ final double[] muArr = new double[originalDim];
for (int c = 0; c < C; c++) {
VectorFloat> mu = landmarks[c];
- // materialize μ once
for (int i = 0; i < originalDim; i++) {
muArr[i] = mu.get(i);
}
- float sum = 0f;
+ double totalSum = 0.0;
float[] proj = landmarkProj[c];
for (int j = 0; j < quantizedDim; j++) {
- float v = vecUtil.ashDotRow(A[j], muArr);
- proj[j] = v;
- sum += v;
+ double rowDot = 0.0;
+ double[] Arow = A[j];
+
+ // Double-precision dot product
+ for (int k = 0; k < originalDim; k++) {
+ rowDot += Arow[k] * muArr[k];
+ }
+
+ proj[j] = (float) rowDot;
+ totalSum += rowDot;
}
- landmarkProjSum[c] = sum;
+ landmarkProjSum[c] = (float) totalSum;
}
return new LandmarkProjections(landmarkProj, landmarkProjSum);
@@ -1037,12 +1067,16 @@ public static StiefelTransform runWithoutTraining(int D, int d, Random rng) {
intW info = new intW(0);
// Workspace query for dgesvd
+ double[] a = gData.clone(); // <- preserve original G for the real factorization
double[] workQuery = new double[1];
- LAPACK.getInstance().dgesvd("S", "S", D, d, gData, D, s, u, D, vt, d, workQuery, -1, info);
+ LAPACK.getInstance().dgesvd("S", "S", D, d, a, D, s, u, D, vt, d, workQuery, -1, info);
+ if (info.val != 0) throw new RuntimeException("SVD workspace query failed info=" + info.val);
+
double[] work = new double[(int) workQuery[0]];
+ a = gData.clone(); // <- ensure the real call sees the original Gaussian matrix
// Execute dgesvd with Job "S" (Thin SVD)
- LAPACK.getInstance().dgesvd("S", "S", D, d, gData, D, s, u, D, vt, d, work, work.length, info);
+ LAPACK.getInstance().dgesvd("S", "S", D, d, a, D, s, u, D, vt, d, work, work.length, info);
if (info.val != 0) {
throw new RuntimeException("Native SVD failed with info=" + info.val);
@@ -1055,21 +1089,6 @@ public static StiefelTransform runWithoutTraining(int D, int d, Random rng) {
RealMatrix W = new Array2DRowRealMatrix(deserializeColumnMajor(wData, D, d), false);
- // Sign Alignment (Crucial for ASH/Random bits)
- // To ensure the "direction" of our bits matches the Gaussian seeds,
- // we flip columns of W so they are positively correlated with G.
- for (int j = 0; j < d; j++) {
- double dot = 0;
- for (int i = 0; i < D; i++) {
- dot += W.getEntry(i, j) * G.getEntry(i, j);
- }
- if (dot < 0.0) {
- for (int i = 0; i < D; i++) {
- W.setEntry(i, j, -W.getEntry(i, j));
- }
- }
- }
-
return new StiefelTransform(W);
}
@@ -1361,10 +1380,10 @@ private static StiefelTransform runItqTrainer(
// PCA basis M = V[:, :d] where X = U S V^T
RealMatrix X = new Array2DRowRealMatrix(xHdNorm, false); // [N×D]
logProgress("\t[stage] Starting Native SVD for PCA...");
- // We only need V, so we tell LAPACK to skip U to save massive amounts of time/memory
- RealMatrix V = computeNativeV(X);
+ // We only need VT, so we tell LAPACK to skip U to save massive amounts of time/memory
+ RealMatrix Vt = computeNativeVt(X);
logProgress("\t[stage] Completed Native SVD...");
- RealMatrix M = V.getSubMatrix(0, D - 1, 0, d - 1); // [D×d]
+ RealMatrix M = Vt.transpose().getSubMatrix(0, D - 1, 0, d - 1); // [D×d]
// Early stopping based on stabilized binary codes (subject to max iterations)
RealMatrix prevXbin = null;
@@ -1530,12 +1549,6 @@ private static RealMatrix gaussianMatrix(int rows, int cols, Random rng) {
return new Array2DRowRealMatrix(m, false);
}
- /** Orthonormal factor: U @ V^T from SVD(M). */
-// private static RealMatrix orthogonalize(RealMatrix m) {
-// SingularValueDecomposition svd = new SingularValueDecomposition(m);
-// return svd.getU().multiply(svd.getVT());
-// }
-
// Basic dgemm routine.
private static RealMatrix nativeMultiply(RealMatrix A, RealMatrix B) {
int m = A.getRowDimension();
@@ -1564,16 +1577,17 @@ private static RealMatrix nativeMultiply(RealMatrix A, RealMatrix B) {
}
// Uses the dgesvd routine but tells LAPACK to skip the U matrix (for PCA). Saves time....
- private static RealMatrix computeNativeV(RealMatrix x) {
+ private static RealMatrix computeNativeVt(RealMatrix x) {
int rows = x.getRowDimension();
int cols = x.getColumnDimension();
- double[] a = serializeColumnMajor(x);
+ double[] a0 = serializeColumnMajor(x);
double[] s = new double[Math.min(rows, cols)];
double[] vt = new double[cols * cols]; // Right singular vectors
intW info = new intW(0);
- // Query optimal workspace size (standard LAPACK practice)
+ // Query optimal workspace size (standard LAPACK practice); use a copy
+ double[] a = a0.clone();
double[] workQuery = new double[1];
LAPACK.getInstance().dgesvd("N", "A", rows, cols, a, rows, s, null, 1, vt, cols, workQuery, -1, info);
@@ -1581,6 +1595,7 @@ private static RealMatrix computeNativeV(RealMatrix x) {
double[] work = new double[lwork];
// Actual call: "N" = No U matrix, "A" = All columns of V^T
+ a = a0.clone(); // Use a fresh copy again
LAPACK.getInstance().dgesvd("N", "A", rows, cols, a, rows, s, null, 1, vt, cols, work, lwork, info);
if (info.val != 0) {
@@ -1634,7 +1649,7 @@ private static RealMatrix orthogonalize(RealMatrix m) {
int k = Math.min(rows, cols); // The rank for thin SVD
// Move to column-major 1D array
- double[] a = serializeColumnMajor(m);
+ double[] a0 = serializeColumnMajor(m);
// Prepare Thin-SVD buffers
// u: [rows * k] instead of [rows * rows]
@@ -1645,6 +1660,7 @@ private static RealMatrix orthogonalize(RealMatrix m) {
intW info = new intW(0);
// Workspace Query
+ double[] a = a0.clone(); // Make a copy
double[] workQuery = new double[1];
LAPACK.getInstance().dgesvd(
"S", "S", rows, cols, a, rows, s, u, rows, vt, k,
@@ -1655,6 +1671,7 @@ private static RealMatrix orthogonalize(RealMatrix m) {
double[] work = new double[lwork];
// Actual Thin-SVD call (Job "S")
+ a = a0.clone(); // Fresh copy again
LAPACK.getInstance().dgesvd(
"S", "S", rows, cols, a, rows, s, u, rows, vt, k,
work, lwork, info
@@ -1708,7 +1725,7 @@ private static double[][] deserializeColumnMajor(double[] data, int rows, int co
}
/**
* Elementwise sign producing +/-1 (no zeros), consistent with ASH encoding convention:
- * bit=1 iff projection > 0, else bit=0.
+ * bit=1 iff projection > 0, else bit=-1.
*/
private static RealMatrix sign01(RealMatrix m) {
final int r = m.getRowDimension();
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index f4f9a34ba..d85b10dd1 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -164,17 +164,17 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = dimension / 4; // 384; // (dimension / 4) + HEADER_BITS; // (dimension / 4) + HEADER_BITS;
+ int encodedBits = (dimension / 2) + HEADER_BITS;
// Payload must be 64-bit aligned for SIMD
int payloadBits = encodedBits - HEADER_BITS;
-// if ((payloadBits & 63) != 0) {
-// throw new IllegalArgumentException(
-// "ASH payloadBits must be 64-bit aligned for SIMD. " +
-// "Got payloadBits=" + payloadBits +
-// " (encodedBits=" + encodedBits +
-// ", HEADER_BITS=" + HEADER_BITS + ")"
-// );
-// }
+ if ((payloadBits & 63) != 0) {
+ throw new IllegalArgumentException(
+ "ASH payloadBits must be 64-bit aligned for SIMD. " +
+ "Got payloadBits=" + payloadBits +
+ " (encodedBits=" + encodedBits +
+ ", HEADER_BITS=" + HEADER_BITS + ")"
+ );
+ }
System.out.println(
"\toriginalDim = " + dimension +
From 74b06420346ecbd69ae12a0e918d96814ccb3463 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Wed, 14 Jan 2026 22:54:59 +0000
Subject: [PATCH 030/243] Partial implementation of ASH single vector scoring
logic SIMD support. Does not work yet.
---
.../jvector/vector/VectorUtilSupport.java | 36 +++++++++++++++++++
1 file changed, 36 insertions(+)
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 32d1f8357..9a84f8a6b 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
@@ -262,6 +262,42 @@ default float ashDotRow(float[] Arow, float[] x) {
return acc;
}
+ /**
+ * Computes maskedAdd for a single vector stored in packed *by-vector* form.
+ * Returns:
+ * maskedAdd =
+ *
+ * packedBits layout:
+ * packedBits[packedBase + w] is 64-bit word w for this vector.
+ *
+ * Default implementation is scalar bit-walk with a register accumulator.
+ * SIMD backends may override.
+ */
+ default float ashMaskedAddAllWords(
+ float[] tildeQ,
+ int d,
+ long[] packedBits,
+ int packedBase, // ord * nWords (by-vector contiguous)
+ int nWords
+ ) {
+ float maskedAdd = 0f; // register accumulator (critical!)
+
+ for (int w = 0; w < nWords; w++) {
+ long word = packedBits[packedBase + w];
+ int baseDim = w * 64;
+
+ // Walk set bits (fast when sparse; matches your block fallback style)
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+ if (idx < d) maskedAdd += tildeQ[idx];
+ word &= (word - 1);
+ }
+ }
+
+ return maskedAdd;
+ }
+
/**
* Computes maskedAdd for a block slice of vectors stored in packed block-column-major form.
* For each lane in [0, blockLen), writes:
From a4d35b9d6c16633a54f48854475184b55b97dbce Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Wed, 14 Jan 2026 22:57:10 +0000
Subject: [PATCH 031/243] Initial ASH integration into graph index. Scalar
scoring with separated vectors. Recall and perf look acceptable.
---
.../graph/similarity/BuildScoreProvider.java | 110 ++++++
.../quantization/AsymmetricHashing.java | 74 +++-
.../jbellis/jvector/example/DistancesASH.java | 18 +-
.../github/jbellis/jvector/example/Grid.java | 88 ++++-
.../example/diagnostics/ASHScoreDebug.java | 316 ++++++++++++++++++
.../example/util/CompressorParameters.java | 34 ++
.../example/util/MultiFileDatasource.java | 6 +-
.../jvector/example/yaml/Compression.java | 32 ++
jvector-examples/yaml-configs/datasets.yml | 37 +-
jvector-examples/yaml-configs/default.yml | 40 ++-
10 files changed, 685 insertions(+), 70 deletions(-)
create mode 100644 jvector-examples/src/main/java/io/github/jbellis/jvector/example/diagnostics/ASHScoreDebug.java
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 f0b184e67..3f0e67c0f 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
@@ -18,6 +18,8 @@
import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
import io.github.jbellis.jvector.graph.RemappedRandomAccessVectorValues;
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
+import io.github.jbellis.jvector.quantization.ASHVectors;
import io.github.jbellis.jvector.quantization.BQVectors;
import io.github.jbellis.jvector.quantization.PQVectors;
import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
@@ -230,4 +232,112 @@ public SearchScoreProvider diversityProviderFor(int node1) {
}
};
}
+
+ /**
+ * Returns a BSP that performs approximate score comparisons using the given ASHVectors.
+ *
+ * Construction (node-node) scoring: Uses symmetric ASH–ASH dot-product approximation
+ * computed from the encoded binary payload + per-vector header fields (scale, offset).
+ *
+ *
Query (vector-node) scoring: Delegates to ASHVectors' query-time score function
+ * (asymmetric float query → ASH vector).
+ *
+ *
Constraint: This provider requires {@code C=1} (single landmark) and
+ * {@link VectorSimilarityFunction#DOT_PRODUCT}.
+ */
+ static BuildScoreProvider ashBuildScoreProvider(VectorSimilarityFunction vsf, ASHVectors ashv) {
+ if (vsf != VectorSimilarityFunction.DOT_PRODUCT) {
+ throw new UnsupportedOperationException("ASH build scoring supports DOT_PRODUCT only");
+ }
+
+ final AsymmetricHashing ash = ashv.getCompressor();
+ if (ash.landmarkCount != 1) {
+ throw new IllegalArgumentException("ASH build scoring requires landmarkCount==1, got " + ash.landmarkCount);
+ }
+
+ final int d = ash.quantizedDim;
+
+ // ||mu||^2 is a constant (C=1) for this scorer.
+ final float muNormSq = VectorUtil.dotProduct(ash.landmarks[0], ash.landmarks[0]);
+
+ return new BuildScoreProvider() {
+ @Override
+ public boolean isExact() {
+ return false;
+ }
+
+ @Override
+ public VectorFloat> approximateCentroid() {
+ // In the C=1 case, ASH uses the dataset mean as its single landmark.
+ // Return a copy so callers can't mutate ASH's internal landmark.
+ return ash.landmarks[0].copy();
+ }
+
+ @Override
+ public SearchScoreProvider searchProviderFor(VectorFloat> vector) {
+ // Query-time asymmetric scoring is already implemented by ASHVectors.
+ return new DefaultSearchScoreProvider(ashv.precomputedScoreFunctionFor(vector, vsf));
+ }
+
+ @Override
+ public SearchScoreProvider searchProviderFor(int node1) {
+
+ final AsymmetricHashing.QuantizedVector v1 = ashv.get(node1);
+
+ return new DefaultSearchScoreProvider(new ScoreFunction() {
+ @Override
+ public boolean isExact() {
+ return false;
+ }
+
+ @Override
+ public float similarityTo(int node2) {
+ final AsymmetricHashing.QuantizedVector v2 = ashv.get(node2);
+
+ final long[] aBits = v1.binaryVector;
+ final long[] bBits = v2.binaryVector;
+
+ int matches = 0;
+
+ int bitBase = 0;
+ for (int w = 0; w < aBits.length && bitBase < d; w++, bitBase += 64) {
+ long aw = aBits[w];
+ long bw = bBits[w];
+
+ // Defensive tail masking (should be unnecessary when d is 64-bit aligned)
+ final int remaining = d - bitBase;
+ if (remaining < 64) {
+ final long mask = (remaining == 64) ? ~0L : ((1L << remaining) - 1L);
+ aw &= mask;
+ bw &= mask;
+ }
+
+ // matches = count of equal bits across positions (XNOR-popcount)
+ matches += Long.bitCount(~(aw ^ bw));
+ }
+
+ // g(.) is sign-coded; with bits interpreted as {0,1} representing { -1, +1 } via s = 2b - 1:
+ // = 2*matches - d
+ final float signDot = 2.0f * matches - (float) d;
+
+ // Reconstruction (C=1):
+ // ≈ ||x-μ|| ||y-μ|| + + - ||μ||^2
+ //
+ // with:
+ // scale = ||x-μ|| / sqrt(d)
+ // offset = - ||μ||^2
+ // and decoder scaling f(z)=d^{-1/2}Az implies ≈ (1/d)
+ // so centered term becomes: scale_x * scale_y *
+ return (v1.scale * v2.scale) * signDot + v1.offset + v2.offset + muNormSq;
+ }
+ });
+ }
+
+ @Override
+ public SearchScoreProvider diversityProviderFor(int node1) {
+ // Keep search/diversity consistent for approximate scoring.
+ return searchProviderFor(node1);
+ }
+ };
+ }
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index 79e0d327e..f4486a17a 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -37,7 +37,6 @@
import java.util.concurrent.ForkJoinPool;
import java.util.Random;
import org.apache.commons.math3.linear.MatrixUtils;
-import org.apache.commons.math3.linear.QRDecomposition;
import org.apache.commons.math3.linear.RealMatrix;
import org.apache.commons.math3.linear.Array2DRowRealMatrix;
import org.apache.commons.math3.linear.SingularValueDecomposition;
@@ -534,22 +533,8 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
landmarks[c] = vectorTypeSupport.readFloatVector(in, dim);
}
- // NOTE: StiefelTransform loading will be added in Step 2.
- // For now, construct a placeholder or re-run the RANDOM initializer.
- StiefelTransform stiefelTransform;
-
- if (optimizer == RANDOM) {
- // Deterministic re-construction for now
- stiefelTransform = runWithoutTraining(
- originalDimension,
- quantizedDim,
- new Random(42)
- );
- } else {
- throw new IllegalArgumentException(
- "ASH optimizer " + optimizer + " requires learned transform; not yet supported in load()"
- );
- }
+ // Load StiefelTransform (required for ALL optimizers, including ITQ/LANDING)
+ StiefelTransform stiefelTransform = readStiefelTransform(in, quantizedDim, originalDimension);
LandmarkProjections lp =
computeLandmarkProjections(
@@ -600,7 +585,9 @@ public void write(DataOutput out, int version) throws IOException {
vectorTypeSupport.writeFloatVector(out, landmarks[c]);
}
- // NOTE: stiefelTransform serialization will be added later
+ // Stiefel transform serialization (required for caching and non-RANDOM optimizers):
+ // Serialize A = W^T in float precision as a [d][D] row-major matrix.
+ writeStiefelTransform(out, stiefelTransform);
}
// ---------------------------------------------------------------------
@@ -1740,6 +1727,57 @@ private static RealMatrix sign01(RealMatrix m) {
return new Array2DRowRealMatrix(out, false);
}
+ private static void writeStiefelTransform(DataOutput out, StiefelTransform st) throws IOException {
+ // Serialize A = W^T in float precision as a [d][D] row-major matrix.
+ // This keeps the format compact and matches the encoding/scoring fast paths.
+ out.writeInt(st.rows);
+ out.writeInt(st.cols);
+
+ for (int i = 0; i < st.rows; i++) {
+ float[] row = st.AFloat[i];
+ if (row.length != st.cols) {
+ throw new IOException("Invalid StiefelTransform AFloat row length: " + row.length + " != " + st.cols);
+ }
+ for (int j = 0; j < st.cols; j++) {
+ out.writeFloat(row[j]);
+ }
+ }
+ }
+
+ private static StiefelTransform readStiefelTransform(RandomAccessReader in, int expectedRows, int expectedCols) throws IOException {
+ int rows = in.readInt();
+ int cols = in.readInt();
+
+ if (rows != expectedRows) {
+ throw new IOException("StiefelTransform rows mismatch: expected " + expectedRows + ", got " + rows);
+ }
+ if (cols != expectedCols) {
+ throw new IOException("StiefelTransform cols mismatch: expected " + expectedCols + ", got " + cols);
+ }
+
+ // Read A = W^T as [rows][cols] (row-major) in float precision.
+ float[][] aFloat = new float[rows][cols];
+ for (int i = 0; i < rows; i++) {
+ float[] row = aFloat[i];
+ for (int j = 0; j < cols; j++) {
+ row[j] = in.readFloat();
+ }
+ }
+
+ // Reconstruct W [D x d] so we can build StiefelTransform(W)
+ // W[j][i] = A[i][j]
+ double[][] wData = new double[cols][rows];
+ for (int i = 0; i < rows; i++) {
+ float[] aRow = aFloat[i];
+ for (int j = 0; j < cols; j++) {
+ wData[j][i] = aRow[j];
+ }
+ }
+
+ RealMatrix W = new Array2DRowRealMatrix(wData, false);
+ return new StiefelTransform(W);
+ }
+
@Override
public int compressorSize() {
// Must exactly match write(out, version)
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index d85b10dd1..0ea4b1456 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -100,8 +100,8 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
final int landmarkCount = 1;
// Define the benchmark size
- int maxQueries = 1_000;
- int maxVectors = 1_000_000;
+ int maxQueries = 10_000;
+ int maxVectors = 10_000_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -164,7 +164,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = (dimension / 2) + HEADER_BITS;
+ int encodedBits = (dimension / 4) + HEADER_BITS;
// Payload must be 64-bit aligned for SIMD
int payloadBits = encodedBits - HEADER_BITS;
if ((payloadBits & 63) != 0) {
@@ -694,12 +694,12 @@ public static void main(String[] args) throws IOException {
// runSIFT();
// runGIST();
// runColbert();
- runCohere100k();
+// runCohere100k();
// runADA();
- runADANoZeros();
- runOpenai1536();
- runOpenai3072();
-// runCap6m();
-// runCohere10m();
+// runADANoZeros();
+// runOpenai1536();
+// runOpenai3072();
+ runCap6m();
+ runCohere10m();
}
}
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 be5d66d43..d47368abf 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
@@ -45,6 +45,10 @@
import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider;
import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
+import io.github.jbellis.jvector.quantization.ASHVectors;
+import io.github.jbellis.jvector.example.diagnostics.ASHScoreDebug;
+
import io.github.jbellis.jvector.quantization.CompressedVectors;
import io.github.jbellis.jvector.quantization.NVQuantization;
import io.github.jbellis.jvector.quantization.PQVectors;
@@ -217,18 +221,60 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List exten
{
var floatVectors = ds.getBaseRavv();
- var pq = (PQVectors) buildCompressor.encodeAll(floatVectors);
- var bsp = BuildScoreProvider.pqBuildScoreProvider(ds.getSimilarityFunction(), pq);
- GraphIndexBuilder builder = new GraphIndexBuilder(bsp, floatVectors.dimension(), M, efConstruction, neighborOverflow, 1.2f, addHierarchy, refineFinalGraph);
+ // Encode once using the build compressor (PQ or ASH)
+ CompressedVectors buildCv = buildCompressor.encodeAll(floatVectors);
+
+ // Optional ASH symmetric-vs-asymmetric debug (off by default)
+ // Enable with: -Djvector.ash.debugSymmetric=true
+ if (Boolean.parseBoolean(System.getProperty("jvector.ash.debugSymmetric", "false"))
+ && buildCv instanceof ASHVectors) {
+ ASHScoreDebug.run(
+ floatVectors,
+ (ASHVectors) buildCv,
+ Integer.parseInt(System.getProperty("jvector.ash.debugSymmetric.pairs", "2000")),
+ Integer.parseInt(System.getProperty("jvector.ash.debugSymmetric.subsetN", "50000")),
+ Integer.parseInt(System.getProperty("jvector.ash.debugSymmetric.topK", "100")),
+ new java.util.Random(123)
+ );
+ }
+
+
+ // PQ-only handle (needed for FUSED_PQ write-time supplier). Null when building with ASH.
+ PQVectors pq = null;
+
+ final BuildScoreProvider bsp;
+ if (buildCv instanceof PQVectors) {
+ pq = (PQVectors) buildCv;
+ bsp = BuildScoreProvider.pqBuildScoreProvider(ds.getSimilarityFunction(), pq);
+ } else if (buildCv instanceof ASHVectors) {
+ bsp = BuildScoreProvider.ashBuildScoreProvider(ds.getSimilarityFunction(), (ASHVectors) buildCv);
+ } else {
+ throw new IllegalArgumentException("Unsupported build compressor output type: " + buildCv.getClass().getName());
+ }
+
+ final PQVectors pqFinal = pq;
+
+ GraphIndexBuilder builder = new GraphIndexBuilder(
+ bsp, floatVectors.dimension(), M, efConstruction, neighborOverflow, 1.2f, addHierarchy, refineFinalGraph
+ );
// use the inline vectors index as the score provider for graph construction
Map, OnDiskGraphIndexWriter> writers = new HashMap<>();
Map, Map>> suppliers = new HashMap<>();
OnDiskGraphIndexWriter scoringWriter = null;
+
int n = 0;
for (var features : featureSets) {
+ // FUSED_PQ requires PQVectors at write time; skip when building with ASH
+ if (features.contains(FeatureId.FUSED_PQ) && pq == null) {
+ System.out.println("Skipping Fused PQ feature because build compressor is not PQ");
+ continue;
+ }
+
var graphPath = testDirectory.resolve("graph" + n++);
- var bws = builderWithSuppliers(features, builder.getGraph(), graphPath, floatVectors, pq.getCompressor());
+ var pqCompressorForFeatures = (pq == null) ? null : pq.getCompressor();
+ var bws = builderWithSuppliers(features, builder.getGraph(), graphPath, floatVectors, pqCompressorForFeatures);
+
var writer = bws.builder.build();
writers.put(features, writer);
suppliers.put(features, bws.suppliers);
@@ -236,6 +282,7 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List exten
scoringWriter = writer;
}
}
+
if (scoringWriter == null) {
throw new IllegalStateException("Bench looks for either NVQ_VECTORS or INLINE_VECTORS feature set for scoring compressed builds.");
}
@@ -262,18 +309,20 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List exten
builder.cleanup();
// write the edge lists and close the writers
- // if our feature set contains Fused PQ, we need a Fused ADC write-time supplier (as we don't have neighbor information during writeInline)
writers.entrySet().stream().parallel().forEach(entry -> {
var writer = entry.getValue();
var features = entry.getKey();
+
Map> writeSuppliers;
if (features.contains(FeatureId.FUSED_PQ)) {
+ // Safe: such features are only present when pq != null (we skipped them otherwise)
writeSuppliers = new EnumMap<>(FeatureId.class);
var view = builder.getGraph().getView();
- writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pq, ordinal));
+ writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(view, pqFinal, ordinal));
} else {
writeSuppliers = Map.of();
}
+
try {
writer.write(writeSuppliers);
writer.close();
@@ -281,6 +330,7 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List exten
throw new UncheckedIOException(e);
}
});
+
builder.close();
double totalTime = (System.nanoTime() - startTime) / 1_000_000_000.0;
System.out.format("Build and write %s in %ss%n", featureSets, totalTime);
@@ -290,6 +340,9 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List exten
Map, ImmutableGraphIndex> indexes = new HashMap<>();
n = 0;
for (var features : featureSets) {
+ if (features.contains(FeatureId.FUSED_PQ) && pq == null) {
+ continue;
+ }
var graphPath = testDirectory.resolve("graph" + n++);
var index = OnDiskGraphIndex.load(ReaderSupplierFactory.open(graphPath));
indexes.put(features, index);
@@ -636,9 +689,26 @@ private static VectorCompressor> getCompressor(Function 0");
+ if (subsetN <= 0) throw new IllegalArgumentException("subsetN must be > 0");
+ if (topK <= 0) throw new IllegalArgumentException("topK must be > 0");
+
+ AsymmetricHashing ash = ashVectors.getCompressor();
+ if (ash.landmarkCount != 1) {
+ throw new IllegalArgumentException("ASHScoreDebug assumes landmarkCount==1, got " + ash.landmarkCount);
+ }
+
+ int n = floatVectors.size();
+ int D = floatVectors.dimension();
+ int d = ash.quantizedDim;
+
+ int effectiveSubsetN = Math.min(subsetN, n);
+
+ // Constant for C=1 (included for manuscript-faithful reconstruction; does not affect ranking)
+ final float muNormSq = VectorUtil.dotProduct(ash.landmarks[0], ash.landmarks[0]);
+
+ System.out.println("=== ASHScoreDebug ===");
+ System.out.println("N=" + n + " D=" + D + " d=" + d
+ + " subsetN=" + effectiveSubsetN + " pairs=" + numPairs + " topK=" + topK);
+
+ // 1) Bit-math sanity on trivial patterns (checks signDot identity)
+ sanityCheckBinaryMath(d);
+
+ // 2) Pair sampling: GT dot, asymmetric score, symmetric score
+ double[] gt = new double[numPairs];
+ double[] asym = new double[numPairs];
+ double[] sym = new double[numPairs];
+
+ for (int t = 0; t < numPairs; t++) {
+ int i = rng.nextInt(effectiveSubsetN);
+ int j = rng.nextInt(effectiveSubsetN);
+
+ VectorFloat> xi = floatVectors.getVector(i);
+ VectorFloat> xj = floatVectors.getVector(j);
+
+ // Ground truth dot
+ gt[t] = VectorUtil.dotProduct(xi, xj);
+
+ // Asymmetric: query xi vs encoded j
+ ScoreFunction.ApproximateScoreFunction f =
+ ashVectors.scoreFunctionFor(xi, VectorSimilarityFunction.DOT_PRODUCT);
+ asym[t] = f.similarityTo(j);
+
+ // Symmetric: encoded i vs encoded j (manuscript-consistent)
+ sym[t] = symmetricDot(ashVectors.get(i), ashVectors.get(j), d, muNormSq);
+ }
+
+ printCorrelations(gt, asym, sym);
+
+ // 3) Top-K overlap sanity for a few query ords (subset brute force)
+ int[] queries = pickQueries(effectiveSubsetN, rng, 5);
+ for (int qi = 0; qi < queries.length; qi++) {
+ int q = queries[qi];
+ System.out.println();
+ System.out.println("--- TopK overlap for q=" + q
+ + " (subsetN=" + effectiveSubsetN + ", topK=" + topK + ") ---");
+ topKOverlapForQuery(floatVectors, ashVectors, q, effectiveSubsetN, topK, d, muNormSq);
+ }
+
+ System.out.println("=== end ASHScoreDebug ===");
+ }
+
+ /**
+ * Symmetric ASH–ASH dot-product approximation (C=1), consistent with:
+ * - g produces sign vectors in {-1,+1}^d (stored as bits b in {0,1}^d via s = 2b - 1)
+ * - f(z) = d^{-1/2} A z implies ≈ (1/d)
+ *
+ * With stored header:
+ * - scale = ||x-μ|| / sqrt(d)
+ * - offset = - ||μ||^2
+ *
+ * Centered dot term:
+ * ||x-μ|| ||y-μ|| ≈ (scale_x * scale_y) *
+ *
+ * And = 2*matches - d, where matches is XNOR-popcount of bit vectors.
+ *
+ * Full dot:
+ * ≈ scale_x*scale_y*(2*matches - d) + offset_x + offset_y + ||μ||^2
+ */
+ public static float symmetricDot(AsymmetricHashing.QuantizedVector v1,
+ AsymmetricHashing.QuantizedVector v2,
+ int d,
+ float muNormSq) {
+ long[] aBits = v1.binaryVector;
+ long[] bBits = v2.binaryVector;
+
+ int matches = 0;
+
+ int bitBase = 0;
+ for (int w = 0; w < aBits.length && bitBase < d; w++, bitBase += 64) {
+ long aw = aBits[w];
+ long bw = bBits[w];
+
+ // Defensive tail masking (should be unnecessary when d is 64-bit aligned)
+ int remaining = d - bitBase;
+ if (remaining < 64) {
+ long mask = (remaining == 64) ? ~0L : ((1L << remaining) - 1L);
+ aw &= mask;
+ bw &= mask;
+ }
+
+ matches += Long.bitCount(~(aw ^ bw)); // XNOR-popcount
+ }
+
+ // sign-dot in {-1,+1}^d
+ float signDot = 2.0f * (float) matches - (float) d;
+
+ return (v1.scale * v2.scale) * signDot + v1.offset + v2.offset + muNormSq;
+ }
+
+ private static void printCorrelations(double[] gt, double[] asym, double[] sym) {
+ double cGtAsym = pearson(gt, asym);
+ double cGtSym = pearson(gt, sym);
+ double cAsymSym = pearson(asym, sym);
+
+ System.out.printf("Pearson corr: GT vs ASYM = %.4f%n", cGtAsym);
+ System.out.printf("Pearson corr: GT vs SYM = %.4f%n", cGtSym);
+ System.out.printf("Pearson corr: ASYM vs SYM = %.4f%n", cAsymSym);
+
+ System.out.printf("Means: GT=%.6f ASYM=%.6f SYM=%.6f%n", mean(gt), mean(asym), mean(sym));
+ System.out.printf("Stdev: GT=%.6f ASYM=%.6f SYM=%.6f%n", stdev(gt), stdev(asym), stdev(sym));
+ }
+
+ private static double pearson(double[] a, double[] b) {
+ if (a.length != b.length) throw new IllegalArgumentException("length mismatch");
+ int n = a.length;
+
+ double ma = mean(a);
+ double mb = mean(b);
+
+ double num = 0.0;
+ double da = 0.0;
+ double db = 0.0;
+ for (int i = 0; i < n; i++) {
+ double xa = a[i] - ma;
+ double xb = b[i] - mb;
+ num += xa * xb;
+ da += xa * xa;
+ db += xb * xb;
+ }
+ double denom = Math.sqrt(da) * Math.sqrt(db);
+ return denom == 0.0 ? 0.0 : (num / denom);
+ }
+
+ private static double mean(double[] a) {
+ double s = 0.0;
+ for (int i = 0; i < a.length; i++) s += a[i];
+ return s / (double) a.length;
+ }
+
+ private static double stdev(double[] a) {
+ double m = mean(a);
+ double s2 = 0.0;
+ for (int i = 0; i < a.length; i++) {
+ double d = a[i] - m;
+ s2 += d * d;
+ }
+ return Math.sqrt(s2 / (double) Math.max(1, a.length - 1));
+ }
+
+ private static int[] pickQueries(int subsetN, Random rng, int count) {
+ int[] qs = new int[count];
+ for (int i = 0; i < count; i++) qs[i] = rng.nextInt(subsetN);
+ return qs;
+ }
+
+ private static void topKOverlapForQuery(RandomAccessVectorValues floatVectors,
+ ASHVectors ashVectors,
+ int q,
+ int subsetN,
+ int topK,
+ int d,
+ float muNormSq) {
+ VectorFloat> xq = floatVectors.getVector(q);
+
+ // Precompute asymmetric score function once
+ ScoreFunction.ApproximateScoreFunction asymF =
+ ashVectors.scoreFunctionFor(xq, VectorSimilarityFunction.DOT_PRODUCT);
+
+ // Encoded q for symmetric
+ AsymmetricHashing.QuantizedVector qEnc = ashVectors.get(q);
+
+ // Score all candidates in subset (excluding self)
+ Scored[] gt = new Scored[subsetN - 1];
+ Scored[] asym = new Scored[subsetN - 1];
+ Scored[] sym = new Scored[subsetN - 1];
+
+ int idx = 0;
+ for (int j = 0; j < subsetN; j++) {
+ if (j == q) continue;
+
+ VectorFloat> xj = floatVectors.getVector(j);
+
+ double gtScore = VectorUtil.dotProduct(xq, xj);
+ double asymScore = asymF.similarityTo(j);
+ double symScore = (double) symmetricDot(qEnc, ashVectors.get(j), d, muNormSq);
+
+ gt[idx] = new Scored(j, gtScore);
+ asym[idx] = new Scored(j, asymScore);
+ sym[idx] = new Scored(j, symScore);
+ idx++;
+ }
+
+ int[] gtTop = topKIds(gt, topK);
+ int[] asymTop = topKIds(asym, topK);
+ int[] symTop = topKIds(sym, topK);
+
+ System.out.printf("overlap(GT, ASYM) = %.2f%%%n", 100.0 * overlapFraction(gtTop, asymTop));
+ System.out.printf("overlap(GT, SYM) = %.2f%%%n", 100.0 * overlapFraction(gtTop, symTop));
+ System.out.printf("overlap(ASYM, SYM) = %.2f%%%n", 100.0 * overlapFraction(asymTop, symTop));
+ }
+
+ private static int[] topKIds(Scored[] arr, int k) {
+ int kk = Math.min(k, arr.length);
+ Arrays.sort(arr, new Comparator() {
+ @Override
+ public int compare(Scored a, Scored b) {
+ return Double.compare(b.score, a.score); // descending
+ }
+ });
+ int[] out = new int[kk];
+ for (int i = 0; i < kk; i++) out[i] = arr[i].id;
+ Arrays.sort(out);
+ return out;
+ }
+
+ private static double overlapFraction(int[] aSorted, int[] bSorted) {
+ int i = 0, j = 0, hit = 0;
+ while (i < aSorted.length && j < bSorted.length) {
+ int va = aSorted[i];
+ int vb = bSorted[j];
+ if (va == vb) { hit++; i++; j++; }
+ else if (va < vb) i++;
+ else j++;
+ }
+ return (double) hit / (double) aSorted.length;
+ }
+
+ private static final class Scored {
+ final int id;
+ final double score;
+ Scored(int id, double score) { this.id = id; this.score = score; }
+ }
+
+ private static void sanityCheckBinaryMath(int d) {
+ if ((d & 63) != 0) {
+ System.out.println("[sanity] skipped: d not multiple of 64");
+ return;
+ }
+
+ // For sign vectors s in {-1,+1}^d represented as bits b via s = 2b - 1:
+ // zeros·zeros = d, ones·ones = d, zeros·ones = -d
+ float exp00 = (float) d;
+ float exp11 = (float) d;
+ float exp01 = -(float) d;
+
+ System.out.printf("[sanity] expected signDot: 00=%.1f 11=%.1f 01=%.1f%n", exp00, exp11, exp01);
+
+ int words = (d + 63) >>> 6;
+ long[] zeros = new long[words];
+ long[] ones = new long[words];
+ Arrays.fill(ones, ~0L);
+
+ float got00 = computeSignDotFromBits(zeros, zeros, d);
+ float got11 = computeSignDotFromBits(ones, ones, d);
+ float got01 = computeSignDotFromBits(zeros, ones, d);
+
+ System.out.printf("[sanity] computed signDot: 00=%.1f 11=%.1f 01=%.1f%n", got00, got11, got01);
+ }
+
+ private static float computeSignDotFromBits(long[] aBits, long[] bBits, int d) {
+ int matches = 0;
+ for (int w = 0; w < aBits.length; w++) {
+ matches += Long.bitCount(~(aBits[w] ^ bBits[w]));
+ }
+ return 2.0f * (float) matches - (float) d;
+ }
+}
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 f84b69938..d7dd7f4c5 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
@@ -16,11 +16,14 @@
package io.github.jbellis.jvector.example.util;
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
import io.github.jbellis.jvector.quantization.BinaryQuantization;
import io.github.jbellis.jvector.quantization.NVQuantization;
import io.github.jbellis.jvector.quantization.ProductQuantization;
import io.github.jbellis.jvector.quantization.VectorCompressor;
+import java.io.IOException;
+
public abstract class CompressorParameters {
public static final CompressorParameters NONE = new NoCompressionParameters();
@@ -100,5 +103,36 @@ public VectorCompressor> computeCompressor(DataSet ds) {
return null;
}
}
+ public static class ASHParameters extends CompressorParameters {
+ private final int optimizer;
+ private final int encodedBits;
+ private final int landmarkCount;
+
+ public ASHParameters(int optimizer, int encodedBits, int landmarkCount) {
+ this.optimizer = optimizer;
+ this.encodedBits = encodedBits;
+ this.landmarkCount = landmarkCount;
+ }
+
+ @Override
+ public VectorCompressor> computeCompressor(DataSet ds) {
+ try {
+ return AsymmetricHashing.initialize(ds.getBaseRavv(), optimizer, encodedBits, landmarkCount);
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ public String idStringFor(DataSet ds) {
+ return String.format("ASH_%s_opt%d_bits%d_C%d", ds.getName(), optimizer, encodedBits, landmarkCount);
+ }
+
+ // TODO set to true for production after debug is complete.
+ @Override
+ public boolean supportsCaching() {
+ return false;
+ }
+ }
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java
index 8bba2bd88..3eb7c2bbb 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java
@@ -51,7 +51,7 @@ public DataSet load() throws IOException {
var baseVectors = SiftLoader.readFvecs("fvec/" + basePath);
var queryVectors = SiftLoader.readFvecs("fvec/" + queriesPath);
var gtVectors = SiftLoader.readIvecs("fvec/" + groundTruthPath);
- return DataSetUtils.getScrubbedDataSet(name, VectorSimilarityFunction.COSINE, baseVectors, queryVectors, gtVectors);
+ return DataSetUtils.getScrubbedDataSet(name, VectorSimilarityFunction.DOT_PRODUCT, baseVectors, queryVectors, gtVectors);
}
public static Map byName = new HashMap<>() {{
@@ -95,6 +95,10 @@ public DataSet load() throws IOException {
"wikipedia_squad/100k/text-embedding-3-small_1536_100000_base_vectors.fvec",
"wikipedia_squad/100k/text-embedding-3-small_1536_100000_query_vectors_10000.fvec",
"wikipedia_squad/100k/text-embedding-3-small_1536_100000_indices_query_10000.ivec"));
+ put("ada002-100k-no-zeros", new MultiFileDatasource("ada002-100k",
+ "ada-002-no-zeros/ada_002_100000_base_vectors_no_zeros.fvec",
+ "ada-002-no-zeros/ada_002_100000_query_vectors_10000_no_zeros.fvec",
+ "ada-002-no-zeros/ada-002_gt_no_zeros.ivec"));
put("ada002-100k", new MultiFileDatasource("ada002-100k",
"wikipedia_squad/100k/ada_002_100000_base_vectors.fvec",
"wikipedia_squad/100k/ada_002_100000_query_vectors_10000.fvec",
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 ef1013e80..42ee2f633 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,38 @@ public Function getCompressorParameters() {
};
case "BQ":
return ds -> new CompressorParameters.BQParameters();
+ case "ASH":
+ return ds -> {
+ if (parameters == null) {
+ throw new IllegalArgumentException("ASH requires parameters: encodedBits, optimizer, landmarkCount");
+ }
+
+ int encodedBits = Integer.parseInt(parameters.get("encodedBits"));
+ int landmarkCount = Integer.parseInt(parameters.getOrDefault("landmarkCount", "1"));
+
+ // optimizer can be numeric or name (RANDOM/ITQ/LANDING)
+ String optStr = parameters.get("optimizer");
+ if (optStr == null) {
+ throw new IllegalArgumentException("ASH requires 'optimizer' parameter");
+ }
+
+ int optimizer;
+ switch (optStr.trim().toUpperCase()) {
+ case "RANDOM":
+ optimizer = io.github.jbellis.jvector.quantization.AsymmetricHashing.RANDOM;
+ break;
+ case "ITQ":
+ optimizer = io.github.jbellis.jvector.quantization.AsymmetricHashing.ITQ;
+ break;
+ case "LANDING":
+ optimizer = io.github.jbellis.jvector.quantization.AsymmetricHashing.LANDING;
+ break;
+ default:
+ optimizer = Integer.parseInt(optStr);
+ }
+
+ return new CompressorParameters.ASHParameters(optimizer, encodedBits, landmarkCount);
+ };
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 3fd67217e..049ce470b 100644
--- a/jvector-examples/yaml-configs/datasets.yml
+++ b/jvector-examples/yaml-configs/datasets.yml
@@ -1,23 +1,24 @@
neighborhood-watch-100k:
- - ada002-100k
- - cohere-english-v3-100k
- - openai-v3-small-100k
- - gecko-100k
- - openai-v3-large-3072-100k
- - openai-v3-large-1536-100k
- - e5-small-v2-100k
- - e5-base-v2-100k
- - e5-large-v2-100k
+# - ada002-100k
+ - ada002-100k-no-zeros
+# - cohere-english-v3-100k
+# - openai-v3-small-100k
+# - gecko-100k
+# - openai-v3-large-3072-100k
+# - openai-v3-large-1536-100k
+# - e5-small-v2-100k
+# - e5-base-v2-100k
+# - e5-large-v2-100k
neighborhood-watch-1M:
- - ada002-1M
- - colbert-1M
+# - ada002-1M
+# - colbert-1M
ann-benchmarks:
- - glove-25-angular.hdf5
- - glove-50-angular.hdf5
- - lastfm-64-dot.hdf5
- - glove-100-angular.hdf5
- - glove-200-angular.hdf5
- - nytimes-256-angular.hdf5
- - sift-128-euclidean.hdf5
+# - glove-25-angular.hdf5
+# - glove-50-angular.hdf5
+# - lastfm-64-dot.hdf5
+# - glove-100-angular.hdf5
+# - glove-200-angular.hdf5
+# - nytimes-256-angular.hdf5
+# - sift-128-euclidean.hdf5
# - deep-image-96-angular.hdf5 # large files not yet supported
# - gist-960-euclidean.hdf5 # large files not yet supported
\ No newline at end of file
diff --git a/jvector-examples/yaml-configs/default.yml b/jvector-examples/yaml-configs/default.yml
index 74a42f973..42f34b7aa 100644
--- a/jvector-examples/yaml-configs/default.yml
+++ b/jvector-examples/yaml-configs/default.yml
@@ -3,23 +3,28 @@ version: 6
dataset: default
construction:
- outDegree: [32]
- efConstruction: [100]
+ outDegree: [64]
+ efConstruction: [200]
neighborOverflow: [1.2f]
addHierarchy: [Yes]
refineFinalGraph: [Yes]
fusedGraph: [No]
compression:
- - type: PQ
+# - type: PQ
+# parameters:
+# m: 192 # we can either specify the integer m or the integer mFactor. In this case, m will be set to the data dimensionality divided by mFactor
+# # mFactor: 8
+# # k: 256 # optional parameter. By default, k=256
+# centerData: No
+# anisotropicThreshold: -1.0 # optional parameter. By default, anisotropicThreshold=-1 (i.e., no anisotropy)
+ - type: ASH
parameters:
- m: 192 # we can either specify the integer m or the integer mFactor. In this case, m will be set to the data dimensionality divided by mFactor
- # mFactor: 8
- # k: 256 # optional parameter. By default, k=256
- centerData: No
- anisotropicThreshold: -1.0 # optional parameter. By default, anisotropicThreshold=-1 (i.e., no anisotropy)
+ encodedBits: "1544"
+ optimizer: "ITQ"
+ landmarkCount: "1"
reranking:
- - NVQ
- useSavedIndexIfExists: Yes
+ - FP
+ useSavedIndexIfExists: No
search:
topKOverquery:
@@ -27,12 +32,17 @@ search:
100: [1.0, 2.0]
useSearchPruning: [Yes]
compression:
- - type: PQ
+# - type: PQ
+# parameters:
+# m: 192
+# # k: 256 # optional parameter. By default, k=256
+# centerData: No
+# anisotropicThreshold: -1.0 # optional parameter. By default, anisotropicThreshold=-1 (i.e., no anisotropy)
+ - type: ASH
parameters:
- m: 192
- # k: 256 # optional parameter. By default, k=256
- centerData: No
- anisotropicThreshold: -1.0 # optional parameter. By default, anisotropicThreshold=-1 (i.e., no anisotropy)
+ encodedBits: "1544"
+ optimizer: "ITQ"
+ landmarkCount: "1"
benchmarks: # full option set, if the whole "benchmarks" section is not specified, a default set will be used
throughput: [ AVG ] # additional options: [AVG, MEDIAN, MAX]
latency: [ AVG ] # additional options: [ AVG, STD, P999 ]
From e9ee89516fc99218266acb3cca136ad76cc56932 Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Sat, 14 Feb 2026 23:50:00 +0000
Subject: [PATCH 032/243] Batch kmeans clustering. Renamed various scorers.
Introduced new packed layout with header interleaving. Introduced
ashMaskedAdd_512.
---
.../jvector/quantization/ASHVectors.java | 233 +++++-
.../quantization/AsymmetricHashing.java | 21 +-
.../MiniBatchKMeansClusterer.java | 419 ++++++++++
.../vector/DefaultVectorUtilSupport.java | 98 +++
.../jvector/vector/VectorUtilSupport.java | 9 +
.../jbellis/jvector/example/DistancesASH.java | 70 +-
.../jbellis/jvector/example/DistancesPQ.java | 19 +-
.../example/util/MultiFileDatasource.java | 10 +-
jvector-examples/yaml-configs/datasets.yml | 8 +-
jvector-examples/yaml-configs/default.yml | 16 +-
.../vector/PanamaVectorUtilSupport.java | 737 ++++++++++++++++++
11 files changed, 1577 insertions(+), 63 deletions(-)
create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MiniBatchKMeansClusterer.java
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index 10e66a487..d69086ce2 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -57,6 +57,19 @@ public class ASHVectors implements CompressedVectors {
final AsymmetricHashing.QuantizedVector[] compressedVectors;
final ASHScorer scorer;
+ enum AshSingleKernel {
+ AUTO,
+ SCALAR,
+ SIMD;
+
+ static AshSingleKernel fromProperty() {
+ String v = System.getProperty("jvector.ash.singleKernel", "auto").toLowerCase();
+ if ("scalar".equals(v)) return SCALAR;
+ if ("simd".equals(v)) return SIMD;
+ return AUTO;
+ }
+ }
+
enum AshBlockKernel {
AUTO,
SCALAR,
@@ -84,6 +97,10 @@ static AshBlockKernel fromProperty() {
private long[] packedBits = null;
private int packedBlockCount = 0;
+ // Flat row-major bits for single-vector SIMD (built lazily)
+ // Layout: [v0_w0, v0_w1, ... v0_wN, v1_w0, ...]
+ private volatile long[] flatPackedVectors = null;
+
/**
* Initialize ASHVectors with an array of ASH-compressed vectors.
*
@@ -192,13 +209,172 @@ public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(
VectorFloat> query,
VectorSimilarityFunction similarityFunction) {
- // Scorer follows ASH paper (single landmark, DOT_PRODUCT only)
- final ASHScorer.ASHScoreFunction f = scorer.scoreFunctionFor(query, similarityFunction);
+ if (similarityFunction != VectorSimilarityFunction.DOT_PRODUCT) {
+ throw new UnsupportedOperationException("ASH scorer supports DOT_PRODUCT only");
+ }
+
+ final AshSingleKernel kernel = AshSingleKernel.fromProperty();
+ final VectorUtilSupport vecUtil = VectorizationProvider.getInstance().getVectorUtilSupport();
+
+ // Scalar forced
+ if (kernel == AshSingleKernel.SCALAR) {
+ final ASHScorer.ASHScoreFunction f = scorer.scoreFunctionFor(query, similarityFunction);
+ return node -> f.similarityTo(compressedVectors[node]);
+ }
- // Wrap QuantizedVector scorer into ordinal-based score function
+ // AUTO/SIMD: use SIMD scorer only when backend supports masked-load kernel;
+ // otherwise fall back to the existing scalar behavior.
+ if ((kernel == AshSingleKernel.SIMD || kernel == AshSingleKernel.AUTO)
+ && vecUtil.supportsAshMaskedLoad()) {
+ final QueryPrecompute qp = precomputeQuery(query);
+ return new SimdASHScoreFunction(qp);
+ }
+
+ // Fallback: existing scalar behavior (stable reference)
+ final ASHScorer.ASHScoreFunction f = scorer.scoreFunctionFor(query, similarityFunction);
return node -> f.similarityTo(compressedVectors[node]);
}
+ // ============================================================================
+ // ASH single vector scoring (Optimized Zero-Copy)
+ // ============================================================================
+
+ private final class SimdASHScoreFunction implements ScoreFunction.ApproximateScoreFunction {
+ private final QueryPrecompute qp;
+ private final VectorUtilSupport vecUtil;
+
+ // Cached Arrays (Zero-Copy)
+ private final float[] tildeQPool;
+ private final long[] allPackedVectors;
+
+ private final int d;
+ private final int words;
+
+ SimdASHScoreFunction(QueryPrecompute qp) {
+ this.qp = qp;
+ this.vecUtil = VectorizationProvider.getInstance().getVectorUtilSupport();
+ this.d = qp.d;
+ this.words = ASHVectors.this.words; // access outer class field
+
+ // 1. Grab direct reference to the pool
+ this.tildeQPool = qp.tildeQPool;
+
+ // 2. Trigger/Get the flat array from the outer class
+ this.allPackedVectors = getFlatPackedVectors();
+ }
+
+ @Override
+ public float similarityTo(int node2) {
+
+ final int c = landmarks[node2] & 0xFF;
+ final int qOffset = c * d;
+ final int packedBase = node2 * words;
+
+ // 2. Run existing optimized v4 Kernel
+ final float maskedAdd = vecUtil.ashMaskedAdd_512(
+ tildeQPool,
+ qOffset,
+ allPackedVectors,
+ packedBase,
+ d,
+ words
+ );
+
+ return scales[node2] * (2f * maskedAdd - qp.sumTildeQ[c]) + qp.dotQMu[c] + offsets[node2];
+ }
+
+// @Override
+// public float similarityTo(int node2) {
+// // A. Identify Landmark (Supports C >= 1)
+// // 'landmarks' is a field in the outer ASHVectors class
+// final int c = landmarks[node2] & 0xFF;
+// final int qOffset = c * d;
+//
+// // B. Identify Vector Location in the flat array
+// // Integer multiply is significantly faster than object pointer chasing
+// final int packedBase = node2 * words;
+//
+// // C. Run Kernel (Zero-Copy)
+// // Note: Ensure your VectorUtilSupport has the 'ashMaskedAddFlat' method signature
+// final float maskedAdd = vecUtil.ashMaskedAddFlatOptimized(
+// tildeQPool,
+// qOffset,
+// allPackedVectors,
+// packedBase,
+// d,
+// words
+// );
+//
+// // D. Linear Combination
+// // scales[] and offsets[] are fields in the outer ASHVectors class
+// return scales[node2] * (2f * maskedAdd - qp.sumTildeQ[c])
+// + qp.dotQMu[c]
+// + offsets[node2];
+// }
+
+ @Override
+ public boolean isExact() {
+ return false;
+ }
+ }
+//
+// // ============================================================================
+// // ASH single vector scoring
+// // ============================================================================
+//
+// private final class SimdASHScoreFunction implements ScoreFunction.ApproximateScoreFunction {
+// private final QueryPrecompute qp;
+// private final VectorUtilSupport vecUtil;
+//
+// private final int d;
+// private final int words;
+//
+// // Scratch for current landmark
+// private final float[] tildeQScratch;
+// private int lastC = -1;
+//
+// SimdASHScoreFunction(QueryPrecompute qp) {
+// this.qp = qp;
+// this.vecUtil = VectorizationProvider.getInstance().getVectorUtilSupport();
+// this.d = qp.d;
+// this.words = AsymmetricHashing.QuantizedVector.wordsForDims(d);
+// this.tildeQScratch = new float[d];
+// }
+//
+// @Override
+// public float similarityTo(int node2) {
+// final int c = landmarks[node2] & 0xFF;
+// final int base = c * d;
+//
+// if (c != lastC) {
+// System.arraycopy(qp.tildeQPool, base, tildeQScratch, 0, d);
+// lastC = c;
+// }
+//
+// // by-vector bits
+// final long[] bits = compressedVectors[node2].binaryVector;
+//
+// // Route masked add through VectorUtilSupport.
+// // Default impl is scalar; Panama can override for SIMD.
+// final float maskedAdd = vecUtil.ashMaskedAddAllWords(
+// tildeQScratch,
+// d,
+// bits,
+// 0, // packedBase (by-vector, so base is 0)
+// words
+// );
+//
+// return scales[node2] * (2f * maskedAdd - qp.sumTildeQ[c])
+// + qp.dotQMu[c]
+// + offsets[node2];
+// }
+//
+// @Override
+// public boolean isExact() {
+// return false;
+// }
+// }
+
// ============================================================================
// ASH block-oriented scoring
// ============================================================================
@@ -216,12 +392,12 @@ public ScoreFunction.ApproximateScoreFunction scoreFunctionFor(
* No block-level optimizations are applied.
*
*/
- private static final class ScalarASHBlockScorer implements ASHBlockScorer {
+ private static final class LegacyASHBlockScorer implements ASHBlockScorer {
private final ASHScorer.ASHScoreFunction scorer;
private final AsymmetricHashing.QuantizedVector[] vectors;
- ScalarASHBlockScorer(
+ LegacyASHBlockScorer(
ASHScorer.ASHScoreFunction scorer,
AsymmetricHashing.QuantizedVector[] vectors) {
this.scorer = scorer;
@@ -268,7 +444,7 @@ public ASHBlockScorer blockScorerFor(
ASHScorer.ASHScoreFunction f =
scorer.scoreFunctionFor(query, similarityFunction);
- return new ScalarASHBlockScorer(f, compressedVectors);
+ return new LegacyASHBlockScorer(f, compressedVectors);
}
/**
@@ -319,7 +495,7 @@ public ASHBlockScorer blockScorerFor(
// Forced scalar path (portable register-accumulator)
// ------------------------------------------------------------
if (kernel == AshBlockKernel.SCALAR) {
- return new BlockScalarASHBlockScorer(
+ return new ScalarASHBlockScorer(
compressedVectors,
scales,
offsets,
@@ -336,7 +512,7 @@ public ASHBlockScorer blockScorerFor(
// or falls back to scalar internally.
ensurePackedBits(blockSize);
- return new BlockSimdASHBlockScorer(
+ return new SimdASHBlockScorer(
compressedVectors,
scales,
offsets,
@@ -375,6 +551,39 @@ private void ensurePackedBits(int blockSize) {
}
}
+ /**
+ * Returns a monolithic long[] containing all quantized bits.
+ *
+ * This method is thread-safe and computes the view lazily.
+ * Use this for single-vector SIMD scoring to avoid pointer chasing.
+ *
+ */
+ public long[] getFlatPackedVectors() {
+ if (flatPackedVectors == null) {
+ synchronized (this) {
+ if (flatPackedVectors == null) {
+ flatPackedVectors = createFlatPackedVectors();
+ }
+ }
+ }
+ return flatPackedVectors;
+ }
+
+ private long[] createFlatPackedVectors() {
+ final int n = compressedVectors.length;
+ if (n == 0) return new long[0];
+
+ // words is already a class field
+ long[] flat = new long[n * words];
+
+ for (int i = 0; i < n; i++) {
+ var v = compressedVectors[i];
+ if (v != null && v.binaryVector != null) {
+ System.arraycopy(v.binaryVector, 0, flat, i * words, words);
+ }
+ }
+ return flat;
+ }
/**
* Block-oriented ASH scorer using register-local accumulation.
@@ -385,7 +594,7 @@ private void ensurePackedBits(int blockSize) {
* into scalar registers.
*
*/
- private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
+ private static final class ScalarASHBlockScorer implements ASHBlockScorer {
private final AsymmetricHashing.QuantizedVector[] vectors;
private final float[] scales;
@@ -398,7 +607,7 @@ private static final class BlockScalarASHBlockScorer implements ASHBlockScorer {
private final int d;
private final int words;
- BlockScalarASHBlockScorer(
+ ScalarASHBlockScorer(
AsymmetricHashing.QuantizedVector[] vectors,
float[] scales,
float[] offsets,
@@ -478,7 +687,7 @@ public void scoreRange(int start, int count, float[] out) {
* avoiding per-lane memory traffic and restoring performance.
*
*/
- private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
+ private static final class SimdASHBlockScorer implements ASHBlockScorer {
private final AsymmetricHashing.QuantizedVector[] vectors;
private final float[] scales;
@@ -506,7 +715,7 @@ private static final class BlockSimdASHBlockScorer implements ASHBlockScorer {
private static final java.util.concurrent.atomic.AtomicBoolean PRINTED_MIXED_ONCE =
new java.util.concurrent.atomic.AtomicBoolean(false);
- BlockSimdASHBlockScorer(
+ SimdASHBlockScorer(
AsymmetricHashing.QuantizedVector[] vectors,
float[] scales,
float[] offsets,
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index f4486a17a..da172ef17 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -310,7 +310,8 @@ public static AsymmetricHashing initialize(RandomAccessVectorValues ravv,
long kmStart = System.nanoTime();
// KMeans++ in JVector returns centroids packed as one big vector [C * D]
- KMeansPlusPlusClusterer km = new KMeansPlusPlusClusterer(points, landmarkCount);
+ // KMeansPlusPlusClusterer km = new KMeansPlusPlusClusterer(points, landmarkCount);
+ MiniBatchKMeansClusterer km = new MiniBatchKMeansClusterer(points, landmarkCount, 1024);
VectorFloat> packed = km.cluster(/*unweightedIterations*/ 20, /*anisotropicIterations*/ 0);
long kmEnd = System.nanoTime();
@@ -398,15 +399,15 @@ private static void validateEncodedBits(
);
}
- if ((payloadBits & 63) != 0) {
- throw new IllegalArgumentException(
- "Invalid encodedBits=" + encodedBits +
- ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
- "for aligned binary payload. " +
- "Got payloadBits=" + payloadBits +
- " (headerBits=" + headerBits + ")."
- );
- }
+// if ((payloadBits & 63) != 0) {
+// throw new IllegalArgumentException(
+// "Invalid encodedBits=" + encodedBits +
+// ". ASH requires (encodedBits - headerBits) to be a multiple of 64 " +
+// "for aligned binary payload. " +
+// "Got payloadBits=" + payloadBits +
+// " (headerBits=" + headerBits + ")."
+// );
+// }
if (payloadBits > originalDim) {
throw new IllegalArgumentException(
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MiniBatchKMeansClusterer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MiniBatchKMeansClusterer.java
new file mode 100644
index 000000000..9a337e980
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MiniBatchKMeansClusterer.java
@@ -0,0 +1,419 @@
+/*
+ * 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.vector.Matrix;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static io.github.jbellis.jvector.util.MathUtil.square;
+import static io.github.jbellis.jvector.vector.VectorUtil.*;
+import static java.lang.Math.max;
+
+/**
+ * A MiniBatch KMeans implementation for float vectors.
+ * Drop-in replacement for KMeansPlusPlusClusterer, optimized for large datasets (10M+).
+ * * Includes:
+ * 1. Mini-Batch Stochastic Updates (fast convergence).
+ * 2. K-Means++ Subsampling Initialization (fast startup).
+ */
+public class MiniBatchKMeansClusterer {
+ private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+ private static final boolean VERBOSE_KMEANS = Boolean.getBoolean("jvector.kmeans.verbose");
+ private static final boolean VERBOSE_KMEANS_TIMING = Boolean.getBoolean("jvector.kmeans.timing");
+
+ public static final float UNWEIGHTED = -1.0f;
+
+ private final int k;
+ private final VectorFloat>[] points;
+ private final int[] assignments;
+ private final VectorFloat> centroids;
+ private final float anisotropicThreshold;
+
+ // MiniBatch specific
+ private final int batchSize;
+ private final int[] centroidCounts;
+
+ /**
+ * @param batchSize Recommended: 1024 or 2048 for large datasets
+ */
+ public MiniBatchKMeansClusterer(VectorFloat>[] points, int k, int batchSize) {
+ this(points, chooseInitialCentroids(points, k), UNWEIGHTED, batchSize);
+ }
+
+ public MiniBatchKMeansClusterer(VectorFloat>[] points, int k, float anisotropicThreshold, int batchSize) {
+ this(points, chooseInitialCentroids(points, k), anisotropicThreshold, batchSize);
+ }
+
+ public MiniBatchKMeansClusterer(VectorFloat>[] points, VectorFloat> centroids, float anisotropicThreshold, int batchSize) {
+ if (Float.isNaN(anisotropicThreshold) || anisotropicThreshold < -1.0 || anisotropicThreshold >= 1.0) {
+ throw new IllegalArgumentException("Valid range for anisotropic threshold T is -1.0 <= t < 1.0");
+ }
+
+ this.points = points;
+ this.k = centroids.length() / points[0].length();
+ this.centroids = centroids.copy();
+ this.anisotropicThreshold = anisotropicThreshold;
+ this.batchSize = Math.min(batchSize, points.length);
+
+ this.centroidCounts = new int[k];
+ this.assignments = new int[points.length];
+ Arrays.fill(this.assignments, -1);
+ }
+
+ static float computeParallelCostMultiplier(double threshold, int dimensions) {
+ assert Double.isFinite(threshold) : "threshold=" + threshold;
+ double parallelCost = threshold * threshold;
+ double perpendicularCost = (1 - parallelCost) / (dimensions - 1);
+ return (float) max(1.0, (parallelCost / perpendicularCost));
+ }
+
+ /**
+ * Performs clustering.
+ * @param unweightedIterations If small (< 100), interpreted as EPOCHS (full passes).
+ * If large, interpreted as TOTAL BATCH UPDATES.
+ * @param anisotropicIterations Number of full-batch anisotropic refinement steps.
+ */
+ public VectorFloat> cluster(int unweightedIterations, int anisotropicIterations) {
+
+ // ------------------------------------------------------------
+ // Mini-Batch Unweighted (isotropic) k-means phase
+ // ------------------------------------------------------------
+ int totalBatches = unweightedIterations;
+ // Interpret small iteration counts as Epochs to ensure sufficient data coverage
+ if (unweightedIterations < 100) {
+ int batchesPerEpoch = Math.max(1, points.length / batchSize);
+ totalBatches = unweightedIterations * batchesPerEpoch;
+ }
+
+ long tTotal0 = System.nanoTime();
+
+ for (int i = 0; i < totalBatches; i++) {
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ clusterOnceMiniBatch();
+
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+ if (VERBOSE_KMEANS_TIMING && i % 100 == 0) {
+ System.out.printf("MiniBatch iter %d time = %.3f ms%n", i, (t1 - t0) * 1e-6);
+ }
+ }
+
+ if (VERBOSE_KMEANS) {
+ System.out.printf("MiniBatch unweighted phase finished. Total batches: %d. Total time: %.3f s%n",
+ totalBatches, (System.nanoTime() - tTotal0) * 1e-9);
+ }
+
+ // ------------------------------------------------------------
+ // Anisotropic refinement phase (Full Batch)
+ // ------------------------------------------------------------
+ if (anisotropicIterations > 0) {
+ // MiniBatch leaves assignments in flux; we must perform a full assignment pass
+ // so the anisotropic phase has valid global clusters.
+ assignAllPoints();
+
+ final int threshold = Math.max(1, (int) (0.01 * points.length));
+
+ for (int i = 0; i < anisotropicIterations; i++) {
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+ int changedCount = clusterOnceAnisotropic();
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ if (VERBOSE_KMEANS) {
+ System.out.printf("KMeans anisotropic iter %d: %d/%d reassigned (%.3f%%)%n",
+ i, changedCount, points.length, 100.0 * changedCount / points.length);
+ }
+
+ if (VERBOSE_KMEANS_TIMING) {
+ System.out.printf(" time = %.3f ms%n", (t1 - t0) * 1e-6);
+ }
+
+ if (changedCount <= threshold) {
+ if (VERBOSE_KMEANS) System.out.println("KMeans anisotropic converged early");
+ break;
+ }
+ }
+ }
+
+ return centroids;
+ }
+
+ /**
+ * Performs a single MiniBatch update step:
+ * 1. Sample batch
+ * 2. Find nearest centroids
+ * 3. Update centroids immediately using learning rate
+ */
+ public void clusterOnceMiniBatch() {
+ var random = ThreadLocalRandom.current();
+
+ // 1. Sample indices
+ int[] batchIndices = new int[batchSize];
+ for (int i = 0; i < batchSize; i++) {
+ batchIndices[i] = random.nextInt(points.length);
+ }
+
+ // 2. Compute assignments for the batch
+ int[] batchAssignments = new int[batchSize];
+ for(int i = 0; i < batchSize; i++) {
+ batchAssignments[i] = getNearestCluster(points[batchIndices[i]]);
+ }
+
+ // 3. Update Centroids (Online)
+ int dim = centroids.length() / k;
+ VectorFloat> tempDiff = vectorTypeSupport.createFloatVector(dim);
+
+ for (int i = 0; i < batchSize; i++) {
+ int pointIdx = batchIndices[i];
+ int clusterIdx = batchAssignments[i];
+ VectorFloat> point = points[pointIdx];
+
+ centroidCounts[clusterIdx]++;
+ float learningRate = 1.0f / centroidCounts[clusterIdx];
+
+ // tempDiff = X
+ tempDiff.copyFrom(point, 0, 0, dim);
+ // tempDiff = X - C_old
+ subFromOffset(tempDiff, centroids, clusterIdx * dim);
+
+ // tempDiff = lr * (X - C_old)
+ scale(tempDiff, learningRate);
+
+ // C_new = C_old + tempDiff
+ // Pass the offset explicitly as the second argument
+ addToOffset(centroids, clusterIdx * dim, tempDiff);
+ }
+ }
+
+ // --- Helpers ---
+
+ private void assignAllPoints() {
+ for (int i = 0; i < points.length; i++) {
+ assignments[i] = getNearestCluster(points[i]);
+ }
+ }
+
+ public int clusterOnceAnisotropic() {
+ updateCentroidsAnisotropic();
+ return updateAssignedPointsAnisotropic();
+ }
+
+ // addToOffset: Adds a SMALL vector (source) into a LARGE vector (dest) at a specific offset.
+ private static void addToOffset(VectorFloat> dest, int destOffset, VectorFloat> source) {
+ // Fix: Iterate over source.length() (Small Temp Vector)
+ for (int i = 0; i < source.length(); i++) {
+ dest.set(destOffset + i, dest.get(destOffset + i) + source.get(i));
+ }
+ }
+
+ // subFromOffset: Subtracts a slice of a LARGE vector (source) from a SMALL vector (dest).
+ // (This logic was actually correct, but verifying strictly ensures safety)
+ private static void subFromOffset(VectorFloat> dest, VectorFloat> source, int sourceOffset) {
+ for (int i = 0; i < dest.length(); i++) {
+ dest.set(i, dest.get(i) - source.get(sourceOffset + i));
+ }
+ }
+
+ // --- Optimized Initialization Logic ---
+
+ /**
+ * Chooses the initial centroids for clustering using K-Means++.
+ * If the dataset is large, uses a random subsample to accelerate initialization.
+ */
+ private static VectorFloat> chooseInitialCentroids(VectorFloat>[] points, int k) {
+ if (k <= 0) throw new IllegalArgumentException("Number of clusters must be positive.");
+ if (k > points.length) throw new IllegalArgumentException("K cannot exceed N");
+
+ // Heuristic: If N is massive, run KMeans++ on a subset only.
+ // Cap at 50,000 points or 10*k to ensure statistical representation without performance penalty.
+ int maxInitSamples = Math.max(k * 10, 50_000);
+
+ if (points.length <= maxInitSamples) {
+ return runStandardKMeansPlusPlus(points, k);
+ } else {
+ if (VERBOSE_KMEANS) {
+ System.out.printf("Subsampling %d points for K-Means++ initialization...%n", maxInitSamples);
+ }
+ var random = ThreadLocalRandom.current();
+ VectorFloat>[] subset = new VectorFloat>[maxInitSamples];
+ for (int i = 0; i < maxInitSamples; i++) {
+ subset[i] = points[random.nextInt(points.length)];
+ }
+ return runStandardKMeansPlusPlus(subset, k);
+ }
+ }
+
+ /**
+ * Standard KMeans++ implementation.
+ * When called by chooseInitialCentroids, 'points' may be the full set or a subset.
+ */
+ private static VectorFloat> runStandardKMeansPlusPlus(VectorFloat>[] points, int k) {
+ var random = ThreadLocalRandom.current();
+ VectorFloat> centroids = vectorTypeSupport.createFloatVector(k * points[0].length());
+
+ float[] distances = new float[points.length];
+ Arrays.fill(distances, Float.MAX_VALUE);
+ VectorFloat> distancesVector = vectorTypeSupport.createFloatVector(distances);
+
+ VectorFloat> firstCentroid = points[random.nextInt(points.length)];
+ centroids.copyFrom(firstCentroid, 0, 0, firstCentroid.length());
+
+ VectorFloat> newDistancesVector = vectorTypeSupport.createFloatVector(points.length);
+ for (int i = 0; i < points.length; i++) {
+ newDistancesVector.set(i, squareL2Distance(points[i], firstCentroid));
+ }
+ VectorUtil.minInPlace(distancesVector, newDistancesVector);
+
+ for (int i = 1; i < k; i++) {
+ float totalDistance = VectorUtil.sum(distancesVector);
+ float r = random.nextFloat() * totalDistance;
+ int selectedIdx = -1;
+
+ for (int j = 0; j < points.length; j++) {
+ r -= distancesVector.get(j);
+ if (r < 1e-6) {
+ selectedIdx = j;
+ break;
+ }
+ }
+ if (selectedIdx == -1) selectedIdx = random.nextInt(points.length);
+
+ VectorFloat> nextCentroid = points[selectedIdx];
+ centroids.copyFrom(nextCentroid, 0, i * nextCentroid.length(), nextCentroid.length());
+
+ for (int j = 0; j < points.length; j++) {
+ newDistancesVector.set(j, squareL2Distance(points[j], nextCentroid));
+ }
+ VectorUtil.minInPlace(distancesVector, newDistancesVector);
+ }
+ return centroids;
+ }
+
+ // --- Standard Distance and Update Logic ---
+
+ private int getNearestCluster(VectorFloat> point) {
+ float minDistance = Float.MAX_VALUE;
+ int nearestCluster = 0;
+ int dim = point.length();
+
+ for (int i = 0; i < k; i++) {
+ float distance = squareL2Distance(point, 0, centroids, i * dim, dim);
+ if (distance < minDistance) {
+ minDistance = distance;
+ nearestCluster = i;
+ }
+ }
+ return nearestCluster;
+ }
+
+ private int updateAssignedPointsAnisotropic() {
+ float pcm = computeParallelCostMultiplier(anisotropicThreshold, points[0].length());
+ int dim = points[0].length();
+
+ float[] cNormSquared = new float[k];
+ for (int i = 0; i < k; i++) {
+ cNormSquared[i] = dotProduct(centroids, i * dim, centroids, i * dim, dim);
+ }
+
+ int changedCount = 0;
+ for (int i = 0; i < points.length; i++) {
+ var x = points[i];
+ var xNormSquared = dotProduct(x, x);
+
+ int index = assignments[i];
+ float minDist = Float.MAX_VALUE;
+ for (int j = 0; j < k; j++) {
+ float dist = weightedDistance(x, j, pcm, cNormSquared[j], xNormSquared);
+ if (dist < minDist) {
+ minDist = dist;
+ index = j;
+ }
+ }
+
+ if (index != assignments[i]) {
+ changedCount++;
+ assignments[i] = index;
+ }
+ }
+ return changedCount;
+ }
+
+ private float weightedDistance(VectorFloat> x, int centroid, float parallelCostMultiplier, float cNormSquared, float xNormSquared) {
+ float cDotX = VectorUtil.dotProduct(centroids, centroid * x.length(), x, 0, x.length());
+ float parallelErrorSubtotal = cDotX - xNormSquared;
+ float residualSquaredNorm = cNormSquared - 2 * cDotX + xNormSquared;
+ float parallelError = square(parallelErrorSubtotal);
+ float perpendicularError = residualSquaredNorm - parallelError;
+ return parallelCostMultiplier * parallelError + perpendicularError;
+ }
+
+ private void updateCentroidsAnisotropic() {
+ int dimensions = points[0].length();
+ float pcm = computeParallelCostMultiplier(anisotropicThreshold, dimensions);
+ float orthogonalCostMultiplier = 1.0f / pcm;
+
+ var pointsByCluster = new HashMap>();
+ for (int i = 0; i < assignments.length; i++) {
+ pointsByCluster.computeIfAbsent(assignments[i], __ -> new ArrayList<>()).add(i);
+ }
+
+ for (int i = 0; i < k; i++) {
+ var L = pointsByCluster.getOrDefault(i, List.of());
+ if (L.isEmpty()) {
+ initializeCentroidToRandomPoint(i);
+ continue;
+ }
+
+ var mean = vectorTypeSupport.createFloatVector(dimensions);
+ var outerProdSums = new Matrix(dimensions, dimensions);
+ for (int j : L) {
+ var point = points[j];
+ // Offset is 0 because 'mean' is a fresh vector of the same dimension as 'point'
+ addToOffset(mean, 0, point);
+ float denom = dotProduct(point, point);
+ if (denom > 0) {
+ var op = Matrix.outerProduct(point, point);
+ op.scale(1.0f / denom);
+ outerProdSums.addInPlace(op);
+ }
+ }
+ outerProdSums.scale((1 - orthogonalCostMultiplier) / L.size());
+ scale(mean, 1.0f / L.size());
+
+ for (int j = 0; j < dimensions; j++) {
+ outerProdSums.addTo(j, j, orthogonalCostMultiplier);
+ }
+
+ var invertedMatrix = outerProdSums.invert();
+ centroids.copyFrom(invertedMatrix.multiply(mean), 0, i * dimensions, dimensions);
+ }
+ }
+
+ private void initializeCentroidToRandomPoint(int i) {
+ var random = ThreadLocalRandom.current();
+ centroids.copyFrom(points[random.nextInt(points.length)], 0, i * points[0].length(), points[0].length());
+ }
+}
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..415bbf0bb 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
@@ -547,4 +547,102 @@ public float nvqUniformLoss(VectorFloat> vector, float minValue, float maxValu
return squaredSum;
}
+ @Override
+ public float ashMaskedAddFlat(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words) {
+ float sum = 0.0f;
+ int baseDim = 0;
+
+ for (int w = 0; w < words; w++) {
+ long word = allPackedVectors[packedBase + w];
+
+ // Standard bit-twiddling iteration (iterate only set bits)
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+
+ // Bounds check required for the last word
+ if (idx < d) {
+ sum += tildeQ[qOffset + idx];
+ }
+
+ // Clear the lowest set bit
+ word &= (word - 1);
+ }
+
+ baseDim += 64;
+ }
+ return sum;
+ }
+
+
+ @Override
+ public float ashMaskedAddFlat_dense(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words) {
+ float sum = 0.0f;
+ int baseDim = 0;
+
+ for (int w = 0; w < words; w++) {
+ long word = allPackedVectors[packedBase + w];
+
+ // Standard bit-twiddling iteration (iterate only set bits)
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+
+ // Bounds check required for the last word
+ if (idx < d) {
+ sum += tildeQ[qOffset + idx];
+ }
+
+ // Clear the lowest set bit
+ word &= (word - 1);
+ }
+
+ baseDim += 64;
+ }
+ return sum;
+ }
+
+ @Override
+ public float ashMaskedAdd_512(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words) {
+ float sum = 0.0f;
+ int baseDim = 0;
+
+ for (int w = 0; w < words; w++) {
+ long word = allPackedVectors[packedBase + w];
+
+ // Standard bit-twiddling iteration (iterate only set bits)
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+
+ // Bounds check required for the last word
+ if (idx < d) {
+ sum += tildeQ[qOffset + idx];
+ }
+
+ // Clear the lowest set bit
+ word &= (word - 1);
+ }
+
+ baseDim += 64;
+ }
+ return sum;
+ }
+
+
}
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 9a84f8a6b..eeaaf223b 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
@@ -298,6 +298,15 @@ default float ashMaskedAddAllWords(
return maskedAdd;
}
+ // In VectorUtilSupport interface
+ float ashMaskedAddFlat(float[] tildeQ, int qOffset, long[] allPackedVectors, int packedBase, int d, int words);
+
+ // In VectorUtilSupport interface
+ float ashMaskedAdd_512(float[] tildeQ, int qOffset, long[] allPackedVectors, int packedBase, int d, int words);
+
+ // In VectorUtilSupport interface
+ float ashMaskedAddFlat_dense(float[] tildeQ, int qOffset, long[] allPackedVectors, int packedBase, int d, int words);
+
/**
* Computes maskedAdd for a block slice of vectors stored in packed block-column-major form.
* For each lane in [0, blockLen), writes:
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
index 0ea4b1456..bd07f7625 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesASH.java
@@ -37,10 +37,23 @@
import java.util.concurrent.ForkJoinTask;
public class DistancesASH {
- private static void printScorerInfo(String label, ASHBlockScorer scorer) {
+
+ static final String blockMode = System.getProperty("jvector.ash.blockKernel", "auto").toLowerCase();
+ static final String singleMode = System.getProperty("jvector.ash.singleKernel", "auto").toLowerCase();
+
+ private static void printScorerInfo(String label, Object scorer) {
+ boolean maskedLoad =
+ io.github.jbellis.jvector.vector.VectorizationProvider.getInstance()
+ .getVectorUtilSupport()
+ .supportsAshMaskedLoad();
+
System.out.println(
"\t[" + label + "] scorer implementation = "
+ scorer.getClass().getName()
+ + " (singleKernel=" + singleMode
+ + ", blockKernel=" + blockMode
+ + ", supportsAshMaskedLoad=" + maskedLoad
+ + ")"
);
}
@@ -76,7 +89,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
Boolean.parseBoolean(System.getProperty("jvector.bench.sanity-check", "false"));
final boolean RUN_RECALL_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.recall", "true"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.recall", "false"));
final int RECALL_K =
Integer.getInteger("jvector.bench.recall.k", 10);
@@ -85,7 +98,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy", "false"));
final boolean RUN_SCALAR_SCORING =
- Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "false"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.scalar-scoring", "true"));
final boolean RUN_FLOAT_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
@@ -164,17 +177,17 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
for (VectorFloat> q : queries) VectorUtil.l2normalize(q);
int dimension = vectors.get(0).length();
- int encodedBits = (dimension / 4) + HEADER_BITS;
+ int encodedBits = 328; // (dimension / 4) + HEADER_BITS;
// Payload must be 64-bit aligned for SIMD
int payloadBits = encodedBits - HEADER_BITS;
- if ((payloadBits & 63) != 0) {
- throw new IllegalArgumentException(
- "ASH payloadBits must be 64-bit aligned for SIMD. " +
- "Got payloadBits=" + payloadBits +
- " (encodedBits=" + encodedBits +
- ", HEADER_BITS=" + HEADER_BITS + ")"
- );
- }
+// if ((payloadBits & 63) != 0) {
+// throw new IllegalArgumentException(
+// "ASH payloadBits must be 64-bit aligned for SIMD. " +
+// "Got payloadBits=" + payloadBits +
+// " (encodedBits=" + encodedBits +
+// ", HEADER_BITS=" + HEADER_BITS + ")"
+// );
+// }
System.out.println(
"\toriginalDim = " + dimension +
@@ -266,6 +279,12 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
newToOldFinal = null;
}
+ {
+ ScoreFunction.ApproximateScoreFunction single =
+ ashVecsFinal.scoreFunctionFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
+ printScorerInfo("single", single);
+ }
+
if(RUN_SANITY_CHECK) {
ASHBlockScorer ref = ashVectors.blockScorerFor(q0, VectorSimilarityFunction.DOT_PRODUCT);
@@ -298,8 +317,8 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
// ------------------------------------------------------------------
// Shared parallelization setup (executor-honoring)
// ------------------------------------------------------------------
- ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool(); // For production
- // ForkJoinPool simdExecutor = new ForkJoinPool(4); // For profiling
+// ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool(); // For production
+ ForkJoinPool simdExecutor = new ForkJoinPool(180); // For profiling
int parallelism = simdExecutor.getParallelism();
int chunkSize = Math.max(1, (finalQueryCount + parallelism - 1) / parallelism);
@@ -438,7 +457,7 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
for (int i = s; i < e; i++) {
VectorFloat> q = finalQueries.get(i);
ScoreFunction.ApproximateScoreFunction f =
- ashVecs.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
+ ashVecsFinal.scoreFunctionFor(q, VectorSimilarityFunction.DOT_PRODUCT);
for (int j = 0; j < finalVectorCount; j++) {
localSum += f.similarityTo(j);
@@ -455,10 +474,25 @@ public static void testASHEncodings(String filenameBase, String filenameQueries,
long ashEnd = System.nanoTime();
- System.out.println("\tScalar ASH dot-product computations took "
- + (ashEnd - ashStart) / 1e9 + " seconds");
+ double singleSeconds = (ashEnd - ashStart) / 1e9;
+ long totalDotProducts = (long) finalQueryCount * (long) finalVectorCount;
+ double singleThroughput = totalDotProducts / singleSeconds;
+
+ System.out.println("\tSingle " + singleMode + " scoring took "
+ + singleSeconds + " seconds");
+
+ System.out.println(
+ "\tSingle " + singleMode + " throughput = "
+ + String.format(java.util.Locale.ROOT, "%.3f", singleThroughput)
+ + " dot-products/sec"
+ + " ("
+ + String.format(java.util.Locale.ROOT, "%.3f", singleThroughput / 1e6)
+ + " Mdot/s)"
+ );
+
System.out.println("\tdummyAccumulator = " + (float) (ashDummy));
System.out.println("--");
+
}
// ==================================================================
@@ -699,7 +733,7 @@ public static void main(String[] args) throws IOException {
// runADANoZeros();
// runOpenai1536();
// runOpenai3072();
- runCap6m();
+// runCap6m();
runCohere10m();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
index df415d326..2448a158f 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/DistancesPQ.java
@@ -52,13 +52,13 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
// Benchmark configuration (runtime flags)
// ------------------------------------------------------------
final boolean RUN_ACCURACY_CHECK =
- Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy-check", "true"));
+ Boolean.parseBoolean(System.getProperty("jvector.bench.accuracy-check", "false"));
final boolean RUN_FLOAT_SCORING =
Boolean.parseBoolean(System.getProperty("jvector.bench.float-scoring", "false"));
// Define the benchmark size
- int maxQueries = 1_000;
- int maxVectors = 100_000;
+ int maxQueries = 10_000;
+ int maxVectors = 10_000_000;
int queryCountInFile = SiftLoader.countFvecs(filenameQueries);
int vectorCountInFile = SiftLoader.countFvecs(filenameBase);
@@ -116,7 +116,8 @@ public static void testPQEncodings(String filenameBase, String filenameQueries)
// ------------------------------------------------------------
// Parallel executor (identical to DistancesASH)
// ------------------------------------------------------------
- ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool();
+// ForkJoinPool simdExecutor = PhysicalCoreExecutor.pool(); // For production
+ ForkJoinPool simdExecutor = new ForkJoinPool(180); // For profiling
int parallelism = simdExecutor.getParallelism();
int chunkSize = Math.max(1, (finalQueryCount + parallelism - 1) / parallelism);
@@ -378,12 +379,12 @@ public static void runCohere10m() throws IOException {
}
public static void main(String[] args) throws IOException {
- runCohere100k();
+// runCohere100k();
// runADA();
- runADANoZeros();
- runOpenai1536();
- runOpenai3072();
+// runADANoZeros();
+// runOpenai1536();
+// runOpenai3072();
// runCap6m();
-// runCohere10m();
+ runCohere10m();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java
index 3eb7c2bbb..a48e5c379 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/MultiFileDatasource.java
@@ -140,8 +140,12 @@ public DataSet load() throws IOException {
DATASET_HASH + "/cap/Caselaw_gte-Qwen2-1.5B_embeddings_query_10k_norm_shuffle.fvecs",
DATASET_HASH + "/cap/cap_1m_gt_norm_shuffle_ip_k100.ivecs"));
put("cap-6M", new MultiFileDatasource("cap-6M",
- DATASET_HASH + "/cap/Caselaw_gte-Qwen2-1.5B_embeddings_base_6m_norm_shuffle.fvecs",
- DATASET_HASH + "/cap/Caselaw_gte-Qwen2-1.5B_embeddings_query_10k_norm_shuffle.fvecs",
- DATASET_HASH + "/cap/cap_6m_gt_norm_shuffle_ip_k100.ivecs"));
+ "cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_base_6m_norm_shuffle.fvecs",
+ "cap-6m/Caselaw_gte-Qwen2-1.5B_embeddings_query_10k_norm_shuffle.fvecs",
+ "cap-6m/cap_6m_gt_norm_shuffle_ip_k100.ivecs"));
+ put("cohere-10M", new MultiFileDatasource("cohere-10M",
+ "cohere-10m/cohere_wiki_en_flat_base_10m_norm.fvecs",
+ "cohere-10m/cohere_wiki_en_flat_query_10k_norm.fvecs",
+ "cohere-10m/cohere_wiki_en_flat_gt_10m_ip_k100.ivecs"));
}};
}
diff --git a/jvector-examples/yaml-configs/datasets.yml b/jvector-examples/yaml-configs/datasets.yml
index 049ce470b..6d2223f24 100644
--- a/jvector-examples/yaml-configs/datasets.yml
+++ b/jvector-examples/yaml-configs/datasets.yml
@@ -1,8 +1,10 @@
neighborhood-watch-100k:
# - ada002-100k
- - ada002-100k-no-zeros
-# - cohere-english-v3-100k
-# - openai-v3-small-100k
+# - ada002-100k-no-zeros
+# - cap-6M
+# - cohere-10M
+ - cohere-english-v3-100k
+ - openai-v3-small-100k
# - gecko-100k
# - openai-v3-large-3072-100k
# - openai-v3-large-1536-100k
diff --git a/jvector-examples/yaml-configs/default.yml b/jvector-examples/yaml-configs/default.yml
index 42f34b7aa..678a3de7e 100644
--- a/jvector-examples/yaml-configs/default.yml
+++ b/jvector-examples/yaml-configs/default.yml
@@ -12,15 +12,15 @@ construction:
compression:
# - type: PQ
# parameters:
-# m: 192 # we can either specify the integer m or the integer mFactor. In this case, m will be set to the data dimensionality divided by mFactor
+# m: 97 # we can either specify the integer m or the integer mFactor. In this case, m will be set to the data dimensionality divided by mFactor
# # mFactor: 8
# # k: 256 # optional parameter. By default, k=256
# centerData: No
# anisotropicThreshold: -1.0 # optional parameter. By default, anisotropicThreshold=-1 (i.e., no anisotropy)
- type: ASH
parameters:
- encodedBits: "1544"
- optimizer: "ITQ"
+ encodedBits: "776"
+ optimizer: "random" # random, ITQ
landmarkCount: "1"
reranking:
- FP
@@ -28,20 +28,20 @@ construction:
search:
topKOverquery:
- 10: [1.0, 2.0, 5.0, 10.0]
- 100: [1.0, 2.0]
+ 10: [1.0]
+# 100: [1.0, 2.0]
useSearchPruning: [Yes]
compression:
# - type: PQ
# parameters:
-# m: 192
+# m: 97
# # k: 256 # optional parameter. By default, k=256
# centerData: No
# anisotropicThreshold: -1.0 # optional parameter. By default, anisotropicThreshold=-1 (i.e., no anisotropy)
- type: ASH
parameters:
- encodedBits: "1544"
- optimizer: "ITQ"
+ encodedBits: "776"
+ optimizer: "random"
landmarkCount: "1"
benchmarks: # full option set, if the whole "benchmarks" section is not specified, a default set will be used
throughput: [ AVG ] # additional options: [AVG, MEDIAN, MAX]
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 d65364ea9..a5bf33353 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
@@ -26,6 +26,7 @@
import jdk.incubator.vector.VectorMask;
import jdk.incubator.vector.VectorOperators;
import jdk.incubator.vector.VectorSpecies;
+import java.util.Objects;
import java.util.List;
@@ -1578,6 +1579,742 @@ public boolean supportsAshMaskedLoad() {
return true;
}
+// // This version uses masked loads.
+// @Override
+// public float ashMaskedAddAllWords(float[] tildeQ,
+// int d,
+// long[] packedBits,
+// int packedBase,
+// int words) {
+// final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+// final int LANES = SPEC.length();
+//
+// FloatVector accVec = FloatVector.zero(SPEC);
+// float scalarAcc = 0f;
+//
+// for (int w = 0; w < words; w++) {
+// long wordBits = packedBits[packedBase + w];
+// int baseDim = w * 64;
+//
+// for (int off = 0; off < 64 && (baseDim + off) < d; off += LANES) {
+// int qBase = baseDim + off;
+// VectorMask inRange = SPEC.indexInRange(qBase, d);
+//
+// long maskBits;
+// if (LANES == 16) maskBits = (wordBits >>> off) & 0xFFFFL;
+// else if (LANES == 8) maskBits = (wordBits >>> off) & 0xFFL;
+// else if (LANES == 4) maskBits = (wordBits >>> off) & 0xFL;
+// else {
+// // unusual width fallback
+// long m = wordBits >>> off;
+// float s = 0f;
+// for (int b = 0; b < LANES; b++) {
+// if (((m >>> b) & 1L) != 0L) {
+// int idx = qBase + b;
+// if (idx < d) s += tildeQ[idx];
+// }
+// }
+// scalarAcc += s;
+// continue;
+// }
+//
+// if (maskBits == 0L) continue;
+//
+// VectorMask bitMask = VectorMask.fromLong(SPEC, maskBits).and(inRange);
+// FloatVector masked = FloatVector.fromArray(SPEC, tildeQ, qBase, bitMask);
+// accVec = accVec.add(masked);
+// }
+// }
+//
+// return accVec.reduceLanes(VectorOperators.ADD) + scalarAcc;
+// }
+
+ @Override
+ public float ashMaskedAddFlat(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int LANES = SPEC.length();
+
+ // 4 Accumulators to hide instruction latency (Pipeline depth ~4)
+ FloatVector acc0 = FloatVector.zero(SPEC);
+ FloatVector acc1 = FloatVector.zero(SPEC);
+ FloatVector acc2 = FloatVector.zero(SPEC);
+ FloatVector acc3 = FloatVector.zero(SPEC);
+
+ // Pre-calculate mask extraction constants based on LANES
+ long laneMask;
+ if (LANES == 16) laneMask = 0xFFFFL; // AVX-512
+ else if (LANES == 8) laneMask = 0xFFL; // AVX2
+ else if (LANES == 4) laneMask = 0xFL; // NEON / SSE
+ else laneMask = (1L << LANES) - 1;
+
+ // Safe limit for Fast Path
+ int safeWords = d / 64;
+
+ // --- HOT LOOP ---
+ for (int w = 0; w < safeWords; w++) {
+ long wordBits = allPackedVectors[packedBase + w];
+ if (wordBits == 0) continue; // Branch prediction handles this well
+
+ int baseDim = w * 64;
+
+ // Iterate 64 bits in steps of LANES (e.g., 0, 8, 16, 24... for AVX2)
+ // We unroll manually inside the loop using modulo to pick accumulators
+ // The JIT is smart enough to unroll this loop completely since LANES is constant.
+ int accumulatorIdx = 0;
+
+ for (int off = 0; off < 64; off += LANES) {
+ // Shift and mask
+ long m = (wordBits >>> off) & laneMask;
+
+ if (m != 0) {
+ VectorMask bitMask = VectorMask.fromLong(SPEC, m);
+ FloatVector loaded = FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + off);
+
+ // Round-robin distribution to accumulators to break dependency chains
+ switch (accumulatorIdx & 3) { // (idx % 4)
+ case 0 -> acc0 = acc0.add(loaded, bitMask);
+ case 1 -> acc1 = acc1.add(loaded, bitMask);
+ case 2 -> acc2 = acc2.add(loaded, bitMask);
+ case 3 -> acc3 = acc3.add(loaded, bitMask);
+ }
+ }
+ accumulatorIdx++;
+ }
+ }
+
+ // --- TAIL LOOP (Robust Bounds) ---
+ // Handles the edge case where d is not a multiple of 64
+ for (int w = safeWords; w < words; w++) {
+ long wordBits = allPackedVectors[packedBase + w];
+ if (wordBits == 0) continue;
+
+ int baseDim = w * 64;
+
+ for (int off = 0; off < 64; off += LANES) {
+ int localIdx = baseDim + off;
+ if (localIdx >= d) break;
+
+ long m = (wordBits >>> off) & laneMask;
+ if (m == 0) continue;
+
+ VectorMask bitMask = VectorMask.fromLong(SPEC, m);
+ int absIdx = qOffset + localIdx;
+
+ if (localIdx + LANES <= d) {
+ // Safe unmasked load
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, absIdx), bitMask);
+ } else {
+ // Boundary masked load
+ VectorMask rangeMask = SPEC.indexInRange(absIdx, qOffset + d);
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, absIdx, rangeMask),
+ bitMask.and(rangeMask));
+ }
+ }
+ }
+
+ // Reduce all accumulators
+ return acc0.add(acc1).add(acc2).add(acc3).reduceLanes(VectorOperators.ADD);
+ }
+
+// @Override
+// public float ashMaskedAddFlat(float[] tildeQ,
+// int qOffset,
+// long[] allPackedVectors,
+// int packedBase,
+// int d,
+// int words) {
+// final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+//
+// FloatVector acc0 = FloatVector.zero(SPEC);
+// FloatVector acc1 = FloatVector.zero(SPEC);
+// FloatVector acc2 = FloatVector.zero(SPEC);
+// FloatVector acc3 = FloatVector.zero(SPEC);
+//
+// int safeWords = d / 64;
+//
+// // --- HOT LOOP (Unrolled & Safe) ---
+// for (int w = 0; w < safeWords; w++) {
+// long wordBits = allPackedVectors[packedBase + w];
+// if (wordBits == 0) continue;
+//
+// int baseDim = w * 64;
+//
+// // Unroll 0
+// long m0 = wordBits & 0xFFFFL;
+// if (m0 != 0) {
+// acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim),
+// VectorMask.fromLong(SPEC, m0));
+// }
+//
+// // Unroll 1
+// long m1 = (wordBits >>> 16) & 0xFFFFL;
+// if (m1 != 0) {
+// acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + 16),
+// VectorMask.fromLong(SPEC, m1));
+// }
+//
+// // Unroll 2
+// long m2 = (wordBits >>> 32) & 0xFFFFL;
+// if (m2 != 0) {
+// acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + 32),
+// VectorMask.fromLong(SPEC, m2));
+// }
+//
+// // Unroll 3
+// long m3 = (wordBits >>> 48) & 0xFFFFL;
+// if (m3 != 0) {
+// acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + 48),
+// VectorMask.fromLong(SPEC, m3));
+// }
+// }
+//
+// // --- TAIL LOOP (Robust Bounds) ---
+// for (int w = safeWords; w < words; w++) {
+// long wordBits = allPackedVectors[packedBase + w];
+// if (wordBits == 0) continue;
+//
+// int baseDim = w * 64;
+//
+// for (int off = 0; off < 64; off += 16) {
+// int localIdx = baseDim + off;
+// if (localIdx >= d) break;
+//
+// long maskBits = (wordBits >>> off) & 0xFFFFL;
+// if (maskBits == 0) continue;
+//
+// VectorMask bitMask = VectorMask.fromLong(SPEC, maskBits);
+// int absIdx = qOffset + localIdx;
+//
+// if (localIdx + 16 <= d) {
+// acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, absIdx), bitMask);
+// } else {
+// VectorMask rangeMask = SPEC.indexInRange(absIdx, qOffset + d);
+// acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, absIdx, rangeMask),
+// bitMask.and(rangeMask));
+// }
+// }
+// }
+//
+// return acc0.add(acc1).add(acc2).add(acc3).reduceLanes(VectorOperators.ADD);
+// }
+
+ public float ashMaskedAddFlatOptimizedv1(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int LANES = SPEC.length();
+
+ FloatVector acc0 = FloatVector.zero(SPEC);
+ FloatVector acc1 = FloatVector.zero(SPEC);
+ FloatVector acc2 = FloatVector.zero(SPEC);
+ FloatVector acc3 = FloatVector.zero(SPEC);
+
+ int safeWords = d / 64;
+
+ // --- OPTIMIZATION 1: Specialized AVX-512 Loop (16 Lanes) ---
+ // This covers the vast majority of server workloads.
+ // We unroll the 64-bit word processing into exactly 4 blocks.
+ if (LANES == 16) {
+ for (int w = 0; w < safeWords; w++) {
+ long wordBits = allPackedVectors[packedBase + w];
+ if (wordBits == 0) continue;
+
+ int baseDim = w * 64;
+
+ // Block 0: Bits 0-15
+ // Note: casting to (int) acts as a fast mask for the low bits
+ long m0 = wordBits & 0xFFFF;
+ if (m0 != 0) {
+ VectorMask mask = VectorMask.fromLong(SPEC, m0);
+ FloatVector loaded = FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim);
+ // OPTIMIZATION 2: Blend (Zero-out) + Unmasked Add
+ // "blend" with 0.0f where mask is unset (not())
+ acc0 = acc0.add(loaded.blend(0.0f, mask.not()));
+ }
+
+ // Block 1: Bits 16-31
+ long m1 = (wordBits >>> 16) & 0xFFFF;
+ if (m1 != 0) {
+ VectorMask mask = VectorMask.fromLong(SPEC, m1);
+ FloatVector loaded = FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + 16);
+ acc1 = acc1.add(loaded.blend(0.0f, mask.not()));
+ }
+
+ // Block 2: Bits 32-47
+ long m2 = (wordBits >>> 32) & 0xFFFF;
+ if (m2 != 0) {
+ VectorMask mask = VectorMask.fromLong(SPEC, m2);
+ FloatVector loaded = FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + 32);
+ acc2 = acc2.add(loaded.blend(0.0f, mask.not()));
+ }
+
+ // Block 3: Bits 48-63
+ long m3 = (wordBits >>> 48); // No mask needed for high bits shift
+ if (m3 != 0) {
+ VectorMask mask = VectorMask.fromLong(SPEC, m3);
+ FloatVector loaded = FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + 48);
+ acc3 = acc3.add(loaded.blend(0.0f, mask.not()));
+ }
+ }
+ }
+ // --- OPTIMIZATION 3: Generic / AVX2 Loop ---
+ else {
+ long laneMask = (1L << LANES) - 1;
+ for (int w = 0; w < safeWords; w++) {
+ long wordBits = allPackedVectors[packedBase + w];
+ if (wordBits == 0) continue;
+
+ int baseDim = w * 64;
+ int accumulatorIdx = 0;
+
+ // Manual step loop to avoid modulo in switch
+ for (int off = 0; off < 64; off += LANES) {
+ long m = (wordBits >>> off) & laneMask;
+ if (m != 0) {
+ VectorMask mask = VectorMask.fromLong(SPEC, m);
+ FloatVector loaded = FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + off);
+
+ // Simple Round-Robin without switch/modulo
+ if (accumulatorIdx == 0) acc0 = acc0.add(loaded.blend(0.0f, mask.not()));
+ else if (accumulatorIdx == 1) acc1 = acc1.add(loaded.blend(0.0f, mask.not()));
+ else if (accumulatorIdx == 2) acc2 = acc2.add(loaded.blend(0.0f, mask.not()));
+ else {
+ acc3 = acc3.add(loaded.blend(0.0f, mask.not()));
+ accumulatorIdx = -1; // Reset for next increment
+ }
+ }
+ accumulatorIdx++;
+ }
+ }
+ }
+
+ // --- TAIL LOOP (Keep your original robust logic here) ---
+ // Your original tail logic handles bounds checking correctly.
+ // Since this is the "cold" path, optimization matters less.
+ for (int w = safeWords; w < words; w++) {
+ // ... (Keep your original Tail Loop implementation) ...
+ }
+
+ return acc0.add(acc1).add(acc2).add(acc3).reduceLanes(VectorOperators.ADD);
+ }
+
+ public float ashMaskedAddFlatOptimizedv2(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int LANES = SPEC.length();
+
+ // Expansion to 8 accumulators to saturate 32 ZMM registers and hide memory latency
+ FloatVector acc0 = FloatVector.zero(SPEC);
+ FloatVector acc1 = FloatVector.zero(SPEC);
+ FloatVector acc2 = FloatVector.zero(SPEC);
+ FloatVector acc3 = FloatVector.zero(SPEC);
+ FloatVector acc4 = FloatVector.zero(SPEC);
+ FloatVector acc5 = FloatVector.zero(SPEC);
+ FloatVector acc6 = FloatVector.zero(SPEC);
+ FloatVector acc7 = FloatVector.zero(SPEC);
+
+ int safeWords = d / 64;
+
+ if (LANES == 16) {
+ // Process 2 words at a time (128 floats) to fill the 8-way unroll
+ int w = 0;
+ for (; w <= safeWords - 2; w += 2) {
+ int baseIdx = qOffset + (w << 6);
+
+ // Word 1 (acc0-acc3)
+ long w0 = allPackedVectors[packedBase + w];
+ if (w0 != 0) {
+ long m0 = w0 & 0xFFFF;
+ if (m0 != 0) acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx, VectorMask.fromLong(SPEC, m0)));
+ long m1 = (w0 >>> 16) & 0xFFFF;
+ if (m1 != 0) acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 16, VectorMask.fromLong(SPEC, m1)));
+ long m2 = (w0 >>> 32) & 0xFFFF;
+ if (m2 != 0) acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 32, VectorMask.fromLong(SPEC, m2)));
+ long m3 = (w0 >>> 48);
+ if (m3 != 0) acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 48, VectorMask.fromLong(SPEC, m3)));
+ }
+
+ // Word 2 (acc4-acc7)
+ long w1 = allPackedVectors[packedBase + w + 1];
+ int nextBaseIdx = baseIdx + 64;
+ if (w1 != 0) {
+ long m4 = w1 & 0xFFFF;
+ if (m4 != 0) acc4 = acc4.add(FloatVector.fromArray(SPEC, tildeQ, nextBaseIdx, VectorMask.fromLong(SPEC, m4)));
+ long m5 = (w1 >>> 16) & 0xFFFF;
+ if (m5 != 0) acc5 = acc5.add(FloatVector.fromArray(SPEC, tildeQ, nextBaseIdx + 16, VectorMask.fromLong(SPEC, m5)));
+ long m6 = (w1 >>> 32) & 0xFFFF;
+ if (m6 != 0) acc6 = acc6.add(FloatVector.fromArray(SPEC, tildeQ, nextBaseIdx + 32, VectorMask.fromLong(SPEC, m6)));
+ long m7 = (w1 >>> 48);
+ if (m7 != 0) acc7 = acc7.add(FloatVector.fromArray(SPEC, tildeQ, nextBaseIdx + 48, VectorMask.fromLong(SPEC, m7)));
+ }
+ }
+
+ // Clean up remaining safe word if odd count
+ for (; w < safeWords; w++) {
+ int baseIdx = qOffset + (w << 6);
+ long wordBits = allPackedVectors[packedBase + w];
+ if (wordBits != 0) {
+ long m0 = wordBits & 0xFFFF;
+ if (m0 != 0) acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx, VectorMask.fromLong(SPEC, m0)));
+ long m1 = (wordBits >>> 16) & 0xFFFF;
+ if (m1 != 0) acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 16, VectorMask.fromLong(SPEC, m1)));
+ long m2 = (wordBits >>> 32) & 0xFFFF;
+ if (m2 != 0) acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 32, VectorMask.fromLong(SPEC, m2)));
+ long m3 = (wordBits >>> 48);
+ if (m3 != 0) acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 48, VectorMask.fromLong(SPEC, m3)));
+ }
+ }
+ } else {
+ // Fallback for AVX2/Generic environments
+ long laneMask = (1L << LANES) - 1;
+ for (int w = 0; w < safeWords; w++) {
+ long wordBits = allPackedVectors[packedBase + w];
+ if (wordBits == 0) continue;
+ int baseDim = w * 64;
+ for (int off = 0; off < 64; off += LANES) {
+ long m = (wordBits >>> off) & laneMask;
+ if (m != 0) {
+ // Use masked add directly to avoid blend/not overhead
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, qOffset + baseDim + off), VectorMask.fromLong(SPEC, m));
+ }
+ }
+ }
+ }
+
+ // Tail loop for non-word-aligned dimensions
+ for (int w = safeWords; w < words; w++) {
+ long wordBits = allPackedVectors[packedBase + w];
+ if (wordBits == 0) continue;
+ int baseDim = w * 64;
+ int remaining = d - baseDim;
+ for (int i = 0; i < remaining; i++) {
+ if ((wordBits & (1L << i)) != 0) {
+ // Scalar fallback for the absolute tail
+ // (Could be vectorized with a length-limited mask if performance here is critical)
+ acc0 = acc0.add(FloatVector.broadcast(SPEC, tildeQ[qOffset + baseDim + i]).blend(0, SPEC.indexInRange(0, 1).not()));
+ }
+ }
+ }
+
+ // Horizontal reduction of the 8 accumulators
+ FloatVector res0 = acc0.add(acc1).add(acc2).add(acc3);
+ FloatVector res1 = acc4.add(acc5).add(acc6).add(acc7);
+ return res0.add(res1).reduceLanes(VectorOperators.ADD);
+ }
+
+ public float ashMaskedAddFlatOptimizedv3(float[] tildeQ, int qOffset, long[] allPackedVectors, int packedBase, int d, int words) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+
+ // 6 Accumulators: Optimized to fit in registers without spilling to cache
+ FloatVector acc0 = FloatVector.zero(SPEC), acc1 = FloatVector.zero(SPEC), acc2 = FloatVector.zero(SPEC);
+ FloatVector acc3 = FloatVector.zero(SPEC), acc4 = FloatVector.zero(SPEC), acc5 = FloatVector.zero(SPEC);
+
+ int safeWords = d / 64;
+ for (int w = 0; w < safeWords; w++) {
+ int baseIdx = qOffset + (w << 6);
+ long word = allPackedVectors[packedBase + w];
+
+ // Direct Masked Add: Collapses 'blend' + 'add' into a single AVX-512 instruction
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx), VectorMask.fromLong(SPEC, word & 0xFFFF));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 16), VectorMask.fromLong(SPEC, (word >>> 16) & 0xFFFF));
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 32), VectorMask.fromLong(SPEC, (word >>> 32) & 0xFFFF));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 48), VectorMask.fromLong(SPEC, (word >>> 48)));
+ }
+
+ // Horizontal reduction
+ return acc0.add(acc1).add(acc2).add(acc3).add(acc4).add(acc5).reduceLanes(VectorOperators.ADD);
+ }
+
+ public float ashMaskedAddFlatOptimizedv4(float[] tildeQ, int qOffset, long[] allPackedVectors, int packedBase, int d, int words) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+
+ // STRICT ADHERENCE TO V3 LEARNINGS: 6 Accumulators to prevent register spilling
+ FloatVector acc0 = FloatVector.zero(SPEC), acc1 = FloatVector.zero(SPEC), acc2 = FloatVector.zero(SPEC);
+ FloatVector acc3 = FloatVector.zero(SPEC), acc4 = FloatVector.zero(SPEC), acc5 = FloatVector.zero(SPEC);
+
+ int safeWords = d / 64;
+ int w = 0;
+
+ // --- NEW OPTIMIZATION: 3-Word Unroll ---
+ // This processes 192 dimensions per loop while staying within the 6-accumulator limit.
+ for (; w <= safeWords - 3; w += 3) {
+ int baseIdx = qOffset + (w << 6);
+
+ // Word 0 -> acc0, acc1, acc2, acc3 (Uses 4 sub-blocks)
+ long wd0 = allPackedVectors[packedBase + w];
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx), VectorMask.fromLong(SPEC, wd0 & 0xFFFF));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 16), VectorMask.fromLong(SPEC, (wd0 >>> 16) & 0xFFFF));
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 32), VectorMask.fromLong(SPEC, (wd0 >>> 32) & 0xFFFF));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 48), VectorMask.fromLong(SPEC, (wd0 >>> 48)));
+
+ // Word 1 -> acc4, acc5, acc0, acc1 (Interleaves to reuse registers)
+ long wd1 = allPackedVectors[packedBase + w + 1];
+ int base1 = baseIdx + 64;
+ acc4 = acc4.add(FloatVector.fromArray(SPEC, tildeQ, base1), VectorMask.fromLong(SPEC, wd1 & 0xFFFF));
+ acc5 = acc5.add(FloatVector.fromArray(SPEC, tildeQ, base1 + 16), VectorMask.fromLong(SPEC, (wd1 >>> 16) & 0xFFFF));
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, base1 + 32), VectorMask.fromLong(SPEC, (wd1 >>> 32) & 0xFFFF));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, base1 + 48), VectorMask.fromLong(SPEC, (wd1 >>> 48)));
+
+ // Word 2 -> acc2, acc3, acc4, acc5
+ long wd2 = allPackedVectors[packedBase + w + 2];
+ int base2 = baseIdx + 128;
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, base2), VectorMask.fromLong(SPEC, wd2 & 0xFFFF));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, base2 + 16), VectorMask.fromLong(SPEC, (wd2 >>> 16) & 0xFFFF));
+ acc4 = acc4.add(FloatVector.fromArray(SPEC, tildeQ, base2 + 32), VectorMask.fromLong(SPEC, (wd2 >>> 32) & 0xFFFF));
+ acc5 = acc5.add(FloatVector.fromArray(SPEC, tildeQ, base2 + 48), VectorMask.fromLong(SPEC, (wd2 >>> 48)));
+ }
+
+ // Standard v3 cleanup for remaining words
+ for (; w < safeWords; w++) {
+ int baseIdx = qOffset + (w << 6);
+ long word = allPackedVectors[packedBase + w];
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx), VectorMask.fromLong(SPEC, word & 0xFFFF));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 16), VectorMask.fromLong(SPEC, (word >>> 16) & 0xFFFF));
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 32), VectorMask.fromLong(SPEC, (word >>> 32) & 0xFFFF));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 48), VectorMask.fromLong(SPEC, (word >>> 48)));
+ }
+
+ // Balanced Reduction Tree (Better than linear add)
+
+ FloatVector sum01 = acc0.add(acc1);
+ FloatVector sum23 = acc2.add(acc3);
+ FloatVector sum45 = acc4.add(acc5);
+ return sum01.add(sum23).add(sum45).reduceLanes(VectorOperators.ADD);
+ }
+
+ public float ashMaskedAdd_512(float[] tildeQ, int qOffset, long[] allPackedVectors, int packedBase, int d, int words) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED; // Assumed 512-bit (16 floats)
+ int safeWords = d / 64;
+
+ // 1. Determine Unroll Strategy based on d
+ // We aim for 8 accumulators for large d, but 384 remains at 6 as requested.
+ int numAcc = (d >= 512) ? 8 : (d == 384 ? 6 : (d >= 128 ? 4 : 2));
+
+ // Initialize required accumulators
+ FloatVector acc0 = FloatVector.zero(SPEC), acc1 = FloatVector.zero(SPEC);
+ FloatVector acc2 = null, acc3 = null, acc4 = null, acc5 = null, acc6 = null, acc7 = null;
+
+ if (numAcc >= 4) { acc2 = FloatVector.zero(SPEC); acc3 = FloatVector.zero(SPEC); }
+ if (numAcc >= 6) { acc4 = FloatVector.zero(SPEC); acc5 = FloatVector.zero(SPEC); }
+ if (numAcc >= 8) { acc6 = FloatVector.zero(SPEC); acc7 = FloatVector.zero(SPEC); }
+
+ int w = 0;
+
+ // 2. Main Unrolled Loop
+ // For d=384, this uses your 3-word unroll.
+ // For d >= 512, we use a 2-word unroll across 8 registers for better throughput.
+ int unrollFactor = (d == 384) ? 3 : (d >= 512 ? 2 : 1);
+
+ for (; w <= safeWords - unrollFactor; w += unrollFactor) {
+ int baseIdx = qOffset + (w << 6);
+
+ // Word 0 Processing
+ long wd0 = allPackedVectors[packedBase + w];
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx), VectorMask.fromLong(SPEC, wd0 & 0xFFFF));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 16), VectorMask.fromLong(SPEC, (wd0 >>> 16) & 0xFFFF));
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 32), VectorMask.fromLong(SPEC, (wd0 >>> 32) & 0xFFFF));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 48), VectorMask.fromLong(SPEC, (wd0 >>> 48)));
+
+ if (unrollFactor >= 2) {
+ long wd1 = allPackedVectors[packedBase + w + 1];
+ int b1 = baseIdx + 64;
+ if (numAcc == 6) { // Special Case for d=384 Interleave
+ acc4 = acc4.add(FloatVector.fromArray(SPEC, tildeQ, b1), VectorMask.fromLong(SPEC, wd1 & 0xFFFF));
+ acc5 = acc5.add(FloatVector.fromArray(SPEC, tildeQ, b1 + 16), VectorMask.fromLong(SPEC, (wd1 >>> 16) & 0xFFFF));
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, b1 + 32), VectorMask.fromLong(SPEC, (wd1 >>> 32) & 0xFFFF));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, b1 + 48), VectorMask.fromLong(SPEC, (wd1 >>> 48)));
+ } else { // Standard 8-accumulator path
+ acc4 = acc4.add(FloatVector.fromArray(SPEC, tildeQ, b1), VectorMask.fromLong(SPEC, wd1 & 0xFFFF));
+ acc5 = acc5.add(FloatVector.fromArray(SPEC, tildeQ, b1 + 16), VectorMask.fromLong(SPEC, (wd1 >>> 16) & 0xFFFF));
+ acc6 = acc6.add(FloatVector.fromArray(SPEC, tildeQ, b1 + 32), VectorMask.fromLong(SPEC, (wd1 >>> 32) & 0xFFFF));
+ acc7 = acc7.add(FloatVector.fromArray(SPEC, tildeQ, b1 + 48), VectorMask.fromLong(SPEC, (wd1 >>> 48)));
+ }
+ }
+
+ if (unrollFactor == 3) { // Tail of the d=384 case
+ long wd2 = allPackedVectors[packedBase + w + 2];
+ int b2 = baseIdx + 128;
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, b2), VectorMask.fromLong(SPEC, wd2 & 0xFFFF));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, b2 + 16), VectorMask.fromLong(SPEC, (wd2 >>> 16) & 0xFFFF));
+ acc4 = acc4.add(FloatVector.fromArray(SPEC, tildeQ, b2 + 32), VectorMask.fromLong(SPEC, (wd2 >>> 32) & 0xFFFF));
+ acc5 = acc5.add(FloatVector.fromArray(SPEC, tildeQ, b2 + 48), VectorMask.fromLong(SPEC, (wd2 >>> 48)));
+ }
+ }
+
+ // 3. Cleanup Loop
+ for (; w < safeWords; w++) {
+ int baseIdx = qOffset + (w << 6);
+ long wd = allPackedVectors[packedBase + w];
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx), VectorMask.fromLong(SPEC, wd & 0xFFFF));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 16), VectorMask.fromLong(SPEC, (wd >>> 16) & 0xFFFF));
+ if (numAcc >= 4) {
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 32), VectorMask.fromLong(SPEC, (wd >>> 32) & 0xFFFF));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, baseIdx + 48), VectorMask.fromLong(SPEC, (wd >>> 48)));
+ }
+ }
+
+ // 4. Dynamic Reduction Tree
+ return performReduction(acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7, numAcc);
+ }
+
+ private float performReduction(FloatVector a0, FloatVector a1, FloatVector a2, FloatVector a3,
+ FloatVector a4, FloatVector a5, FloatVector a6, FloatVector a7, int count) {
+ FloatVector res;
+ if (count == 8) {
+ FloatVector s01 = a0.add(a1), s23 = a2.add(a3), s45 = a4.add(a5), s67 = a6.add(a7);
+ res = s01.add(s23).add(s45.add(s67));
+ } else if (count == 6) {
+ res = a0.add(a1).add(a2.add(a3)).add(a4.add(a5));
+ } else if (count == 4) {
+ res = a0.add(a1).add(a2.add(a3));
+ } else {
+ res = a0.add(a1);
+ }
+ return res.reduceLanes(VectorOperators.ADD);
+ }
+
+ // This version is optimized for dense masks and uses a masked load, plain add.
+ // It is branchless and fully unrolled.
+ // It only works for 64-bit aligned payloads and 16 lanes (AVX512).
+ public float ashMaskedAddFlat_dense(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words_not_used) { // TODO words is not used in this version
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int LANES = SPEC.length();
+
+ // Typical case: d is 256/384 => multiple of 64
+ // If you can guarantee that, you can delete all tail handling.
+ final int words = d >>> 6;
+
+ Objects.checkFromIndexSize(qOffset, d, tildeQ.length);
+ Objects.checkFromIndexSize(packedBase, words, allPackedVectors.length);
+
+ // 4 accumulators is fine for AVX-512; keep it.
+ FloatVector acc0 = FloatVector.zero(SPEC);
+ FloatVector acc1 = FloatVector.zero(SPEC);
+ FloatVector acc2 = FloatVector.zero(SPEC);
+ FloatVector acc3 = FloatVector.zero(SPEC);
+
+ if (LANES == 16) {
+ for (int w = 0, qBase = qOffset; w < words; w++, qBase += 64) {
+ long bits = allPackedVectors[packedBase + w];
+
+ // Build masks from shifted bits (no laneMask needed)
+ VectorMask k0 = VectorMask.fromLong(SPEC, bits);
+ VectorMask k1 = VectorMask.fromLong(SPEC, bits >>> 16);
+ VectorMask k2 = VectorMask.fromLong(SPEC, bits >>> 32);
+ VectorMask k3 = VectorMask.fromLong(SPEC, bits >>> 48);
+
+ // Masked load (zeroing) + unmasked add: branchless
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, qBase + 0, k0));
+ acc1 = acc1.add(FloatVector.fromArray(SPEC, tildeQ, qBase + 16, k1));
+ acc2 = acc2.add(FloatVector.fromArray(SPEC, tildeQ, qBase + 32, k2));
+ acc3 = acc3.add(FloatVector.fromArray(SPEC, tildeQ, qBase + 48, k3));
+ }
+ } else {
+ // Generic fallback (still branchless and no laneMask)
+ for (int w = 0, qBase = qOffset; w < words; w++, qBase += 64) {
+ long bits = allPackedVectors[packedBase + w];
+ for (int off = 0; off < 64; off += LANES) {
+ VectorMask k = VectorMask.fromLong(SPEC, bits >>> off);
+ acc0 = acc0.add(FloatVector.fromArray(SPEC, tildeQ, qBase + off, k));
+ }
+ }
+ }
+
+ return acc0.add(acc1).add(acc2).add(acc3).reduceLanes(VectorOperators.ADD);
+ }
+
+ // This approach uses masked adds instead of masked loads, and two accumulators (since this is slow)
+ @Override
+ public float ashMaskedAddAllWords(float[] tildeQ,
+ int d,
+ long[] packedBits,
+ int packedBase,
+ int words) {
+ final VectorSpecies SPEC = FloatVector.SPECIES_PREFERRED;
+ final int LANES = SPEC.length();
+ final int upperBoundSafe = d - LANES;
+
+ // Two accumulators to hide floating-point addition latency (4 cycles)
+ FloatVector acc1 = FloatVector.zero(SPEC);
+ FloatVector acc2 = FloatVector.zero(SPEC);
+ float scalarAcc = 0f;
+
+ for (int w = 0; w < words; w++) {
+ long wordBits = packedBits[packedBase + w];
+ // REMOVED: if (wordBits == 0) continue; (useless for dense data)
+
+ int baseDim = w * 64;
+
+ // Toggle to alternate between acc1 and acc2
+ boolean toggle = false;
+
+ for (int off = 0; off < 64; off += LANES) {
+ int qBase = baseDim + off;
+
+ // Keep the fast integer check!
+ if (qBase >= d) break;
+
+ long maskBits;
+ if (LANES == 16) maskBits = (wordBits >>> off) & 0xFFFFL;
+ else if (LANES == 8) maskBits = (wordBits >>> off) & 0xFFL;
+ else if (LANES == 4) maskBits = (wordBits >>> off) & 0xFL;
+ else {
+ // Fallback for unusual widths
+ long m = wordBits >>> off;
+ for (int b = 0; b < LANES; b++) {
+ if (((m >>> b) & 1L) != 0L) {
+ int idx = qBase + b;
+ if (idx < d) scalarAcc += tildeQ[idx];
+ }
+ }
+ continue;
+ }
+
+ // Still worth checking if the local chunk is 0 to skip the ADD
+ if (maskBits == 0L) continue;
+
+ VectorMask bitMask = VectorMask.fromLong(SPEC, maskBits);
+ FloatVector loaded;
+
+ // Fast Unmasked Load (Still critical for dense data)
+ if (qBase <= upperBoundSafe) {
+ loaded = FloatVector.fromArray(SPEC, tildeQ, qBase);
+ } else {
+ VectorMask rangeMask = SPEC.indexInRange(qBase, d);
+ loaded = FloatVector.fromArray(SPEC, tildeQ, qBase, bitMask.and(rangeMask));
+ }
+
+ // Alternate accumulators to parallelize the ADDs
+ if (!toggle) {
+ acc1 = acc1.add(loaded, bitMask);
+ } else {
+ acc2 = acc2.add(loaded, bitMask);
+ }
+ toggle = !toggle;
+ }
+ }
+
+ // Merge the two pipelines
+ return acc1.add(acc2).reduceLanes(VectorOperators.ADD) + scalarAcc;
+ }
+
@Override
public void ashMaskedAddBlockAllWords(float[] tildeQ,
int d,
From 2f4e02d5d7b89a62f53ecce555da341703ba4c4a Mon Sep 17 00:00:00 2001
From: Ted Willke
Date: Sun, 15 Feb 2026 01:29:00 +0000
Subject: [PATCH 033/243] Additional cleanup for the merge of latest main into
branch.
---
.../jvector/quantization/ASHVectors.java | 4 ++--
.../quantization/AsymmetricHashing.java | 12 +++++------
.../benchmarks/datasets/DataSetLoaderMFD.java | 2 +-
.../jvector/example/util/SiftLoader.java | 4 +++-
jvector-examples/yaml-configs/default.yml | 20 +++++++++----------
5 files changed, 22 insertions(+), 20 deletions(-)
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
index d69086ce2..9759322ad 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ASHVectors.java
@@ -16,6 +16,7 @@
package io.github.jbellis.jvector.quantization;
+import io.github.jbellis.jvector.disk.IndexWriter;
import io.github.jbellis.jvector.disk.RandomAccessReader;
import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
import io.github.jbellis.jvector.util.RamUsageEstimator;
@@ -25,7 +26,6 @@
import io.github.jbellis.jvector.vector.VectorizationProvider;
import io.github.jbellis.jvector.vector.types.VectorFloat;
-import java.io.DataOutput;
import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
@@ -174,7 +174,7 @@ public int count() {
* carefully when introduced.
*/
@Override
- public void write(DataOutput out, int version) throws IOException {
+ public void write(IndexWriter out, int version) throws IOException {
// Write ASH compressor first
ash.write(out, version);
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
index da172ef17..06960e69d 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/AsymmetricHashing.java
@@ -17,6 +17,7 @@
package io.github.jbellis.jvector.quantization;
import io.github.jbellis.jvector.annotations.VisibleForTesting;
+import io.github.jbellis.jvector.disk.IndexWriter;
import io.github.jbellis.jvector.disk.RandomAccessReader;
import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
@@ -30,7 +31,6 @@
import dev.ludovic.netlib.blas.BLAS;
import dev.ludovic.netlib.lapack.LAPACK;
-import java.io.DataOutput;
import java.io.IOException;
import java.util.Arrays;
import java.util.Objects;
@@ -561,11 +561,11 @@ public static AsymmetricHashing load(RandomAccessReader in) throws IOException {
/**
* TODO confirm handling versioning correctly
* TODO confirm handling magic number correctly
- * Writes the ASH index header to DataOutput (used for on-disk indexing)
- * @param out the DataOutput into which to write the object
+ * Writes the ASH index header to IndexWriter (used for on-disk indexing)
+ * @param out the IndexWriter into which to write the object
* @throws IOException if there is a problem writing to out.
*/
- public void write(DataOutput out, int version) throws IOException {
+ public void write(IndexWriter out, int version) throws IOException {
if (version > OnDiskGraphIndex.CURRENT_VERSION) {
throw new IllegalArgumentException("Unsupported version " + version);
}
@@ -826,7 +826,7 @@ static void quantizeTo(VectorFloat> vector,
);
}
- public void write(DataOutput out, int quantizedDim) throws IOException {
+ public void write(IndexWriter out, int quantizedDim) throws IOException {
out.writeFloat(scale);
out.writeFloat(offset);
out.writeByte(landmark);
@@ -1728,7 +1728,7 @@ private static RealMatrix sign01(RealMatrix m) {
return new Array2DRowRealMatrix(out, false);
}
- private static void writeStiefelTransform(DataOutput out, StiefelTransform st) throws IOException {
+ private static void writeStiefelTransform(IndexWriter out, StiefelTransform st) throws IOException {
// Serialize A = W^T in float precision as a [d][D] row-major matrix.
// This keeps the format compact and matches the encoding/scoring fast paths.
out.writeInt(st.rows);
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderMFD.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderMFD.java
index b20c0db50..769c1cfae 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderMFD.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetLoaderMFD.java
@@ -152,7 +152,7 @@ public DataSet load() {
var baseVectors = SiftLoader.readFvecs("fvec/" + basePath);
var queryVectors = SiftLoader.readFvecs("fvec/" + queriesPath);
var gtVectors = SiftLoader.readIvecs("fvec/" + groundTruthPath);
- return DataSetUtils.getScrubbedDataSet(name, VectorSimilarityFunction.COSINE, baseVectors, queryVectors, gtVectors);
+ return DataSetUtils.getScrubbedDataSet(name, VectorSimilarityFunction.DOT_PRODUCT, baseVectors, queryVectors, gtVectors);
}
public static Map byName = new HashMap<>() {{
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 f68ae340d..24d5f9e81 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
@@ -46,7 +46,7 @@ public class SiftLoader {
* Reads the entire file and materializes all vectors into memory.
*