diff --git a/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/PackAsBinaryBenchmark.java b/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/PackAsBinaryBenchmark.java index babfcfcc84745..38ab305a01441 100644 --- a/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/PackAsBinaryBenchmark.java +++ b/benchmarks/src/main/java/org/elasticsearch/benchmark/vector/PackAsBinaryBenchmark.java @@ -82,4 +82,13 @@ public void packAsBinaryLegacy(Blackhole bh) { bh.consume(packed); } } + + @Benchmark + @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" }) + public void packAsBinaryPanama(Blackhole bh) { + for (int i = 0; i < numVectors; i++) { + BQVectorUtils.packAsBinary(qVectors[i], packed); + bh.consume(packed); + } + } } diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java index a02370a89f931..81fa6a959574f 100644 --- a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java +++ b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/ESVectorUtil.java @@ -368,4 +368,17 @@ public static void soarDistanceBulk( } IMPL.soarDistanceBulk(v1, c0, c1, c2, c3, originalResidual, soarLambda, rnorm, distances); } + + /** + * Packs the provided int array populated with "0" and "1" values into a byte array. + * + * @param vector the int array to pack, must contain only "0" and "1" values. + * @param packed the byte array to store the packed result, must be large enough to hold the packed data. + */ + public static void packAsBinary(int[] vector, byte[] packed) { + if (packed.length * Byte.SIZE < vector.length) { + throw new IllegalArgumentException("packed array is too small: " + packed.length * Byte.SIZE + " < " + vector.length); + } + IMPL.packAsBinary(vector, packed); + } } diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java index 12abda2506252..f9c52a544db9c 100644 --- a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java +++ b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/DefaultESVectorUtilSupport.java @@ -320,4 +320,37 @@ public void soarDistanceBulk( distances[2] = soarDistance(v1, c2, originalResidual, soarLambda, rnorm); distances[3] = soarDistance(v1, c3, originalResidual, soarLambda, rnorm); } + + @Override + public void packAsBinary(int[] vector, byte[] packed) { + packAsBinaryImpl(vector, packed); + } + + public static void packAsBinaryImpl(int[] vector, byte[] packed) { + int limit = vector.length - 7; + int i = 0; + int index = 0; + for (; i < limit; i += 8, index++) { + assert vector[i] == 0 || vector[i] == 1; + assert vector[i + 1] == 0 || vector[i + 1] == 1; + assert vector[i + 2] == 0 || vector[i + 2] == 1; + assert vector[i + 3] == 0 || vector[i + 3] == 1; + assert vector[i + 4] == 0 || vector[i + 4] == 1; + assert vector[i + 5] == 0 || vector[i + 5] == 1; + assert vector[i + 6] == 0 || vector[i + 6] == 1; + assert vector[i + 7] == 0 || vector[i + 7] == 1; + int result = vector[i] << 7 | (vector[i + 1] << 6) | (vector[i + 2] << 5) | (vector[i + 3] << 4) | (vector[i + 4] << 3) + | (vector[i + 5] << 2) | (vector[i + 6] << 1) | (vector[i + 7]); + packed[index] = (byte) result; + } + if (i == vector.length) { + return; + } + byte result = 0; + for (int j = 7; j >= 0 && i < vector.length; i++, j--) { + assert vector[i] == 0 || vector[i] == 1; + result |= (byte) ((vector[i] & 1) << j); + } + packed[index] = result; + } } diff --git a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java index 895105a452b0c..d75b9cf6cfd25 100644 --- a/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java +++ b/libs/simdvec/src/main/java/org/elasticsearch/simdvec/internal/vectorization/ESVectorUtilSupport.java @@ -63,4 +63,6 @@ void soarDistanceBulk( float rnorm, float[] distances ); + + void packAsBinary(int[] vector, byte[] packed); } diff --git a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java index 2a5f633d51b78..33c2eca4c6152 100644 --- a/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java +++ b/libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorUtilSupport.java @@ -22,8 +22,10 @@ import org.apache.lucene.util.Constants; import static jdk.incubator.vector.VectorOperators.ADD; +import static jdk.incubator.vector.VectorOperators.LSHL; import static jdk.incubator.vector.VectorOperators.MAX; import static jdk.incubator.vector.VectorOperators.MIN; +import static jdk.incubator.vector.VectorOperators.OR; public final class PanamaESVectorUtilSupport implements ESVectorUtilSupport { @@ -942,4 +944,81 @@ public void soarDistanceBulk( distances[2] = dsq2 + soarLambda * proj2 * proj2 / rnorm; distances[3] = dsq3 + soarLambda * proj3 * proj3 / rnorm; } + + private static final VectorSpecies INT_SPECIES_128 = IntVector.SPECIES_128; + private static final IntVector SHIFTS_256; + private static final IntVector HIGH_SHIFTS_128; + private static final IntVector LOW_SHIFTS_128; + static { + final int[] shifts = new int[] { 7, 6, 5, 4, 3, 2, 1, 0 }; + if (VECTOR_BITSIZE == 128) { + HIGH_SHIFTS_128 = IntVector.fromArray(INT_SPECIES_128, shifts, 0); + LOW_SHIFTS_128 = IntVector.fromArray(INT_SPECIES_128, shifts, INT_SPECIES_128.length()); + SHIFTS_256 = null; + } else { + SHIFTS_256 = IntVector.fromArray(INT_SPECIES_256, shifts, 0); + HIGH_SHIFTS_128 = null; + LOW_SHIFTS_128 = null; + } + } + private static final int[] SHIFTS = new int[] { 7, 6, 5, 4, 3, 2, 1, 0 }; + + @Override + public void packAsBinary(int[] vector, byte[] packed) { + // 128 / 32 == 4 + if (vector.length >= 8 && HAS_FAST_INTEGER_VECTORS) { + // TODO: can we optimize for >= 512? + if (VECTOR_BITSIZE >= 256) { + packAsBinary256(vector, packed); + return; + } else if (VECTOR_BITSIZE == 128) { + packAsBinary128(vector, packed); + return; + } + } + DefaultESVectorUtilSupport.packAsBinaryImpl(vector, packed); + } + + private void packAsBinary256(int[] vector, byte[] packed) { + final int limit = INT_SPECIES_256.loopBound(vector.length); + int i = 0; + int index = 0; + for (; i < limit; i += INT_SPECIES_256.length(), index++) { + IntVector v = IntVector.fromArray(INT_SPECIES_256, vector, i); + int result = v.lanewise(LSHL, SHIFTS_256).reduceLanes(OR); + packed[index] = (byte) result; + } + if (i == vector.length) { + return; // all done + } + byte result = 0; + for (int j = 7; j >= 0 && i < vector.length; i++, j--) { + assert vector[i] == 0 || vector[i] == 1; + result |= (byte) ((vector[i] & 1) << j); + } + packed[index] = result; + } + + private void packAsBinary128(int[] vector, byte[] packed) { + final int limit = INT_SPECIES_128.loopBound(vector.length) - INT_SPECIES_128.length(); + int i = 0; + int index = 0; + for (; i < limit; i += 2 * INT_SPECIES_128.length(), index++) { + IntVector v = IntVector.fromArray(INT_SPECIES_128, vector, i); + var v1 = v.lanewise(LSHL, HIGH_SHIFTS_128); + v = IntVector.fromArray(INT_SPECIES_128, vector, i + INT_SPECIES_128.length()); + var v2 = v.lanewise(LSHL, LOW_SHIFTS_128); + int result = v1.lanewise(OR, v2).reduceLanes(OR); + packed[index] = (byte) result; + } + if (i == vector.length) { + return; // all done + } + byte result = 0; + for (int j = 7; j >= 0 && i < vector.length; i++, j--) { + assert vector[i] == 0 || vector[i] == 1; + result |= (byte) ((vector[i] & 1) << j); + } + packed[index] = result; + } } diff --git a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java index b51fc25fab9f1..b0aa4a9a45afe 100644 --- a/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java +++ b/libs/simdvec/src/test/java/org/elasticsearch/simdvec/ESVectorUtilTests.java @@ -9,6 +9,7 @@ package org.elasticsearch.simdvec; +import org.elasticsearch.index.codec.vectors.BQVectorUtils; import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer; import org.elasticsearch.simdvec.internal.vectorization.BaseVectorizationTests; import org.elasticsearch.simdvec.internal.vectorization.ESVectorizationProvider; @@ -355,6 +356,20 @@ public void testSoarDistanceBulk() { assertArrayEquals(expectedDistances, panamaDistances, deltaEps); } + public void testPackAsBinary() { + int dims = randomIntBetween(16, 2048); + int[] toPack = new int[dims]; + for (int i = 0; i < dims; i++) { + toPack[i] = randomInt(1); + } + int length = BQVectorUtils.discretize(dims, 64) / 8; + byte[] packed = new byte[length]; + byte[] packedLegacy = new byte[length]; + defaultedProvider.getVectorUtilSupport().packAsBinary(toPack, packedLegacy); + defOrPanamaProvider.getVectorUtilSupport().packAsBinary(toPack, packed); + assertArrayEquals(packedLegacy, packed); + } + private float[] generateRandomVector(int size) { float[] vector = new float[size]; for (int i = 0; i < size; ++i) { diff --git a/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java b/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java index f2ef2b05541f8..cba55f8a7e942 100644 --- a/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java +++ b/server/src/main/java/org/elasticsearch/index/codec/vectors/BQVectorUtils.java @@ -22,6 +22,7 @@ import org.apache.lucene.util.ArrayUtil; import org.apache.lucene.util.BitUtil; import org.apache.lucene.util.VectorUtil; +import org.elasticsearch.simdvec.ESVectorUtil; /** Utility class for vector quantization calculations */ public class BQVectorUtils { @@ -55,31 +56,7 @@ public static void packAsBinaryLegacy(int[] vector, byte[] packed) { } public static void packAsBinary(int[] vector, byte[] packed) { - int limit = vector.length - 7; - int i = 0; - int index = 0; - for (; i < limit; i += 8, index++) { - assert vector[i] == 0 || vector[i] == 1; - assert vector[i + 1] == 0 || vector[i + 1] == 1; - assert vector[i + 2] == 0 || vector[i + 2] == 1; - assert vector[i + 3] == 0 || vector[i + 3] == 1; - assert vector[i + 4] == 0 || vector[i + 4] == 1; - assert vector[i + 5] == 0 || vector[i + 5] == 1; - assert vector[i + 6] == 0 || vector[i + 6] == 1; - assert vector[i + 7] == 0 || vector[i + 7] == 1; - int result = vector[i] << 7 | (vector[i + 1] << 6) | (vector[i + 2] << 5) | (vector[i + 3] << 4) | (vector[i + 4] << 3) - | (vector[i + 5] << 2) | (vector[i + 6] << 1) | (vector[i + 7]); - packed[index] = (byte) result; - } - if (i == vector.length) { - return; - } - byte result = 0; - for (int j = 7; j >= 0 && i < vector.length; i++, j--) { - assert vector[i] == 0 || vector[i] == 1; - result |= (byte) ((vector[i] & 1) << j); - } - packed[index] = result; + ESVectorUtil.packAsBinary(vector, packed); } public static int discretize(int value, int bucket) {