You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by us...@apache.org on 2023/10/15 18:01:50 UTC
[lucene] 03/04: speedup all binary functions on avx256, speedup binary square on avx512 (#12681)
This is an automated email from the ASF dual-hosted git repository.
uschindler pushed a commit to branch branch_9x
in repository https://gitbox.apache.org/repos/asf/lucene.git
commit ebaa34a3065c41a23de16f49a85293f96f9dd49e
Author: Robert Muir <rm...@apache.org>
AuthorDate: Sat Oct 14 19:08:36 2023 +0200
speedup all binary functions on avx256, speedup binary square on avx512 (#12681)
- Move testing properties to provider class (no classloading deadlock possible) and fallback to default provider in non-test mode
- fix html
- tidy
- fix Panama provider to only have static final constants, but allow sysprops for testing.
- Set those sysprops by the build when running tests: force integer vectors on always, and randomize vector size
- stop using SPECIES_PREFERRED except at the top of this file
- speedup all binary functions on avx256, speedup binary square on avx512
- cleanup cosine too, no perf impact
- simple cleanups to vector code
---
gradle/testing/randomization.gradle | 7 +
.../vectorization/VectorizationProvider.java | 52 +-
.../vectorization/PanamaVectorUtilSupport.java | 736 +++++++++++++--------
.../vectorization/PanamaVectorizationProvider.java | 36 +-
4 files changed, 505 insertions(+), 326 deletions(-)
diff --git a/gradle/testing/randomization.gradle b/gradle/testing/randomization.gradle
index e4b395b1019..9b49a53bcf0 100644
--- a/gradle/testing/randomization.gradle
+++ b/gradle/testing/randomization.gradle
@@ -105,6 +105,13 @@ allprojects {
// miscellaneous; some of them very weird.
[propName: 'tests.LUCENE_VERSION', value: baseVersion, description: "Base Lucene version."],
[propName: 'tests.bwcdir', value: null, description: "Data for backward-compatibility indexes."],
+ // vectorization related
+ [propName: 'tests.vectorsize',
+ value: { ->
+ RandomPicks.randomFrom(new Random(projectSeedLong), ["128", "256", "512"])
+ },
+ description: "Sets preferred vector size in bits."],
+ [propName: 'tests.forceintegervectors', value: "true", description: "Forces use of integer vectors even when slow."],
]
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java b/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java
index 9f0040f3e28..ed4066a94ac 100644
--- a/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java
+++ b/lucene/core/src/java/org/apache/lucene/internal/vectorization/VectorizationProvider.java
@@ -26,8 +26,11 @@ import java.security.PrivilegedAction;
import java.util.Locale;
import java.util.Objects;
import java.util.Optional;
+import java.util.OptionalInt;
import java.util.Set;
+import java.util.function.Predicate;
import java.util.logging.Logger;
+import java.util.stream.Stream;
import org.apache.lucene.util.SuppressForbidden;
import org.apache.lucene.util.VectorUtil;
@@ -40,6 +43,35 @@ import org.apache.lucene.util.VectorUtil;
*/
public abstract class VectorizationProvider {
+ static final OptionalInt TESTS_VECTOR_SIZE;
+ static final boolean TESTS_FORCE_INTEGER_VECTORS;
+
+ static {
+ var vs = OptionalInt.empty();
+ try {
+ vs =
+ Stream.ofNullable(System.getProperty("tests.vectorsize"))
+ .filter(Predicate.not(String::isEmpty))
+ .mapToInt(Integer::parseInt)
+ .findAny();
+ } catch (
+ @SuppressWarnings("unused")
+ SecurityException se) {
+ // ignored
+ }
+ TESTS_VECTOR_SIZE = vs;
+
+ boolean enforce = false;
+ try {
+ enforce = Boolean.getBoolean("tests.forceintegervectors");
+ } catch (
+ @SuppressWarnings("unused")
+ SecurityException se) {
+ // ignored
+ }
+ TESTS_FORCE_INTEGER_VECTORS = enforce;
+ }
+
/**
* Returns the default instance of the provider matching vectorization possibilities of actual
* runtime.
@@ -90,10 +122,17 @@ public abstract class VectorizationProvider {
return new DefaultVectorizationProvider();
}
vectorMod.ifPresent(VectorizationProvider.class.getModule()::addReads);
- // check if client VM
- if (!testMode && isClientVM()) {
- LOG.warning("C2 compiler is disabled; Java vector incubator API can't be enabled");
- return new DefaultVectorizationProvider();
+ // check for testMode and otherwise fallback to default if slowness could happen
+ if (!testMode) {
+ if (TESTS_VECTOR_SIZE.isPresent() || TESTS_FORCE_INTEGER_VECTORS) {
+ LOG.warning(
+ "Vector bitsize and/or integer vectors enforcement; using default vectorization provider outside of testMode");
+ return new DefaultVectorizationProvider();
+ }
+ if (isClientVM()) {
+ LOG.warning("C2 compiler is disabled; Java vector incubator API can't be enabled");
+ return new DefaultVectorizationProvider();
+ }
}
try {
// we use method handles with lookup, so we do not need to deal with setAccessible as we
@@ -102,10 +141,9 @@ public abstract class VectorizationProvider {
final var cls =
lookup.findClass(
"org.apache.lucene.internal.vectorization.PanamaVectorizationProvider");
- final var constr =
- lookup.findConstructor(cls, MethodType.methodType(void.class, boolean.class));
+ final var constr = lookup.findConstructor(cls, MethodType.methodType(void.class));
try {
- return (VectorizationProvider) constr.invoke(testMode);
+ return (VectorizationProvider) constr.invoke();
} catch (UnsupportedOperationException uoe) {
// not supported because preferred vector size too small or similar
LOG.warning("Java vector incubator API was not enabled. " + uoe.getMessage());
diff --git a/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java b/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java
index c066304136f..e509f6ef49c 100644
--- a/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java
+++ b/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java
@@ -16,165 +16,145 @@
*/
package org.apache.lucene.internal.vectorization;
+import static jdk.incubator.vector.VectorOperators.ADD;
+import static jdk.incubator.vector.VectorOperators.B2I;
+import static jdk.incubator.vector.VectorOperators.B2S;
+import static jdk.incubator.vector.VectorOperators.S2I;
+
import jdk.incubator.vector.ByteVector;
import jdk.incubator.vector.FloatVector;
import jdk.incubator.vector.IntVector;
import jdk.incubator.vector.ShortVector;
import jdk.incubator.vector.Vector;
-import jdk.incubator.vector.VectorOperators;
import jdk.incubator.vector.VectorShape;
import jdk.incubator.vector.VectorSpecies;
+import org.apache.lucene.util.Constants;
+/**
+ * VectorUtil methods implemented with Panama incubating vector API.
+ *
+ * <p>Supports two system properties for correctness testing purposes only:
+ *
+ * <ul>
+ * <li>tests.vectorsize (int)
+ * <li>tests.forceintegervectors (boolean)
+ * </ul>
+ *
+ * Setting these properties will make this code run EXTREMELY slow!
+ */
final class PanamaVectorUtilSupport implements VectorUtilSupport {
- private static final int INT_SPECIES_PREF_BIT_SIZE = IntVector.SPECIES_PREFERRED.vectorBitSize();
+ // preferred vector sizes, which can be altered for testing
+ private static final VectorSpecies<Float> FLOAT_SPECIES;
+ private static final VectorSpecies<Integer> INT_SPECIES;
+ private static final VectorSpecies<Byte> BYTE_SPECIES;
+ private static final VectorSpecies<Short> SHORT_SPECIES;
- private static final VectorSpecies<Float> PREF_FLOAT_SPECIES = FloatVector.SPECIES_PREFERRED;
- private static final VectorSpecies<Byte> PREF_BYTE_SPECIES;
- private static final VectorSpecies<Short> PREF_SHORT_SPECIES;
+ static final int VECTOR_BITSIZE;
+ static final boolean HAS_FAST_INTEGER_VECTORS;
static {
- if (INT_SPECIES_PREF_BIT_SIZE >= 256) {
- PREF_BYTE_SPECIES =
- ByteVector.SPECIES_MAX.withShape(
- VectorShape.forBitSize(IntVector.SPECIES_PREFERRED.vectorBitSize() >> 2));
- PREF_SHORT_SPECIES =
- ShortVector.SPECIES_MAX.withShape(
- VectorShape.forBitSize(IntVector.SPECIES_PREFERRED.vectorBitSize() >> 1));
+ // default to platform supported bitsize
+ int vectorBitSize = VectorShape.preferredShape().vectorBitSize();
+ // but allow easy overriding for testing
+ vectorBitSize = VectorizationProvider.TESTS_VECTOR_SIZE.orElse(vectorBitSize);
+ INT_SPECIES = VectorSpecies.of(int.class, VectorShape.forBitSize(vectorBitSize));
+ VECTOR_BITSIZE = INT_SPECIES.vectorBitSize();
+ FLOAT_SPECIES = INT_SPECIES.withLanes(float.class);
+ // compute BYTE/SHORT sizes relative to preferred integer vector size
+ if (VECTOR_BITSIZE >= 256) {
+ BYTE_SPECIES = ByteVector.SPECIES_MAX.withShape(VectorShape.forBitSize(VECTOR_BITSIZE >> 2));
+ SHORT_SPECIES =
+ ShortVector.SPECIES_MAX.withShape(VectorShape.forBitSize(VECTOR_BITSIZE >> 1));
} else {
- PREF_BYTE_SPECIES = null;
- PREF_SHORT_SPECIES = null;
+ BYTE_SPECIES = null;
+ SHORT_SPECIES = null;
}
- }
-
- private final boolean useIntegerVectors;
-
- PanamaVectorUtilSupport(boolean useIntegerVectors) {
- this.useIntegerVectors = useIntegerVectors;
+ // hotspot misses some SSE intrinsics, workaround it
+ // to be fair, they do document this thing only works well with AVX2/AVX3 and Neon
+ boolean isAMD64withoutAVX2 = Constants.OS_ARCH.equals("amd64") && VECTOR_BITSIZE < 256;
+ HAS_FAST_INTEGER_VECTORS =
+ VectorizationProvider.TESTS_FORCE_INTEGER_VECTORS || (isAMD64withoutAVX2 == false);
}
@Override
public float dotProduct(float[] a, float[] b) {
int i = 0;
float res = 0;
+
// if the array size is large (> 2x platform vector size), its worth the overhead to vectorize
- if (a.length > 2 * PREF_FLOAT_SPECIES.length()) {
- // vector loop is unrolled 4x (4 accumulators in parallel)
- FloatVector acc1 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector acc2 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector acc3 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector acc4 = FloatVector.zero(PREF_FLOAT_SPECIES);
- int upperBound = PREF_FLOAT_SPECIES.loopBound(a.length - 3 * PREF_FLOAT_SPECIES.length());
- for (; i < upperBound; i += 4 * PREF_FLOAT_SPECIES.length()) {
- FloatVector va = FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i);
- FloatVector vb = FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i);
- acc1 = acc1.add(va.mul(vb));
- FloatVector vc =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + PREF_FLOAT_SPECIES.length());
- FloatVector vd =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + PREF_FLOAT_SPECIES.length());
- acc2 = acc2.add(vc.mul(vd));
- FloatVector ve =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + 2 * PREF_FLOAT_SPECIES.length());
- FloatVector vf =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + 2 * PREF_FLOAT_SPECIES.length());
- acc3 = acc3.add(ve.mul(vf));
- FloatVector vg =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + 3 * PREF_FLOAT_SPECIES.length());
- FloatVector vh =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + 3 * PREF_FLOAT_SPECIES.length());
- acc4 = acc4.add(vg.mul(vh));
- }
- // vector tail: less scalar computations for unaligned sizes, esp with big vector sizes
- upperBound = PREF_FLOAT_SPECIES.loopBound(a.length);
- for (; i < upperBound; i += PREF_FLOAT_SPECIES.length()) {
- FloatVector va = FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i);
- FloatVector vb = FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i);
- acc1 = acc1.add(va.mul(vb));
- }
- // reduce
- FloatVector res1 = acc1.add(acc2);
- FloatVector res2 = acc3.add(acc4);
- res += res1.add(res2).reduceLanes(VectorOperators.ADD);
+ if (a.length > 2 * FLOAT_SPECIES.length()) {
+ i += FLOAT_SPECIES.loopBound(a.length);
+ res += dotProductBody(a, b, i);
}
+ // scalar tail
for (; i < a.length; i++) {
res += b[i] * a[i];
}
return res;
}
+ /** vectorized float dot product body */
+ private float dotProductBody(float[] a, float[] b, int limit) {
+ int i = 0;
+ // vector loop is unrolled 4x (4 accumulators in parallel)
+ // we don't know how many the cpu can do at once, some can do 2, some 4
+ FloatVector acc1 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector acc2 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector acc3 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector acc4 = FloatVector.zero(FLOAT_SPECIES);
+ int unrolledLimit = limit - 3 * FLOAT_SPECIES.length();
+ for (; i < unrolledLimit; i += 4 * FLOAT_SPECIES.length()) {
+ // one
+ FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
+ FloatVector vb = FloatVector.fromArray(FLOAT_SPECIES, b, i);
+ acc1 = acc1.add(va.mul(vb));
+
+ // two
+ FloatVector vc = FloatVector.fromArray(FLOAT_SPECIES, a, i + FLOAT_SPECIES.length());
+ FloatVector vd = FloatVector.fromArray(FLOAT_SPECIES, b, i + FLOAT_SPECIES.length());
+ acc2 = acc2.add(vc.mul(vd));
+
+ // three
+ FloatVector ve = FloatVector.fromArray(FLOAT_SPECIES, a, i + 2 * FLOAT_SPECIES.length());
+ FloatVector vf = FloatVector.fromArray(FLOAT_SPECIES, b, i + 2 * FLOAT_SPECIES.length());
+ acc3 = acc3.add(ve.mul(vf));
+
+ // four
+ FloatVector vg = FloatVector.fromArray(FLOAT_SPECIES, a, i + 3 * FLOAT_SPECIES.length());
+ FloatVector vh = FloatVector.fromArray(FLOAT_SPECIES, b, i + 3 * FLOAT_SPECIES.length());
+ acc4 = acc4.add(vg.mul(vh));
+ }
+ // vector tail: less scalar computations for unaligned sizes, esp with big vector sizes
+ for (; i < limit; i += FLOAT_SPECIES.length()) {
+ FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
+ FloatVector vb = FloatVector.fromArray(FLOAT_SPECIES, b, i);
+ acc1 = acc1.add(va.mul(vb));
+ }
+ // reduce
+ FloatVector res1 = acc1.add(acc2);
+ FloatVector res2 = acc3.add(acc4);
+ return res1.add(res2).reduceLanes(ADD);
+ }
+
@Override
public float cosine(float[] a, float[] b) {
int i = 0;
float sum = 0;
float norm1 = 0;
float norm2 = 0;
+
// if the array size is large (> 2x platform vector size), its worth the overhead to vectorize
- if (a.length > 2 * PREF_FLOAT_SPECIES.length()) {
- // vector loop is unrolled 4x (4 accumulators in parallel)
- FloatVector sum1 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector sum2 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector sum3 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector sum4 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm1_1 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm1_2 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm1_3 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm1_4 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm2_1 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm2_2 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm2_3 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector norm2_4 = FloatVector.zero(PREF_FLOAT_SPECIES);
- int upperBound = PREF_FLOAT_SPECIES.loopBound(a.length - 3 * PREF_FLOAT_SPECIES.length());
- for (; i < upperBound; i += 4 * PREF_FLOAT_SPECIES.length()) {
- FloatVector va = FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i);
- FloatVector vb = FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i);
- sum1 = sum1.add(va.mul(vb));
- norm1_1 = norm1_1.add(va.mul(va));
- norm2_1 = norm2_1.add(vb.mul(vb));
- FloatVector vc =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + PREF_FLOAT_SPECIES.length());
- FloatVector vd =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + PREF_FLOAT_SPECIES.length());
- sum2 = sum2.add(vc.mul(vd));
- norm1_2 = norm1_2.add(vc.mul(vc));
- norm2_2 = norm2_2.add(vd.mul(vd));
- FloatVector ve =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + 2 * PREF_FLOAT_SPECIES.length());
- FloatVector vf =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + 2 * PREF_FLOAT_SPECIES.length());
- sum3 = sum3.add(ve.mul(vf));
- norm1_3 = norm1_3.add(ve.mul(ve));
- norm2_3 = norm2_3.add(vf.mul(vf));
- FloatVector vg =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + 3 * PREF_FLOAT_SPECIES.length());
- FloatVector vh =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + 3 * PREF_FLOAT_SPECIES.length());
- sum4 = sum4.add(vg.mul(vh));
- norm1_4 = norm1_4.add(vg.mul(vg));
- norm2_4 = norm2_4.add(vh.mul(vh));
- }
- // vector tail: less scalar computations for unaligned sizes, esp with big vector sizes
- upperBound = PREF_FLOAT_SPECIES.loopBound(a.length);
- for (; i < upperBound; i += PREF_FLOAT_SPECIES.length()) {
- FloatVector va = FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i);
- FloatVector vb = FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i);
- sum1 = sum1.add(va.mul(vb));
- norm1_1 = norm1_1.add(va.mul(va));
- norm2_1 = norm2_1.add(vb.mul(vb));
- }
- // reduce
- FloatVector sumres1 = sum1.add(sum2);
- FloatVector sumres2 = sum3.add(sum4);
- FloatVector norm1res1 = norm1_1.add(norm1_2);
- FloatVector norm1res2 = norm1_3.add(norm1_4);
- FloatVector norm2res1 = norm2_1.add(norm2_2);
- FloatVector norm2res2 = norm2_3.add(norm2_4);
- sum += sumres1.add(sumres2).reduceLanes(VectorOperators.ADD);
- norm1 += norm1res1.add(norm1res2).reduceLanes(VectorOperators.ADD);
- norm2 += norm2res1.add(norm2res2).reduceLanes(VectorOperators.ADD);
+ if (a.length > 2 * FLOAT_SPECIES.length()) {
+ i += FLOAT_SPECIES.loopBound(a.length);
+ float[] ret = cosineBody(a, b, i);
+ sum += ret[0];
+ norm1 += ret[1];
+ norm2 += ret[2];
}
+ // scalar tail
for (; i < a.length; i++) {
float elem1 = a[i];
float elem2 = b[i];
@@ -185,56 +165,87 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
return (float) (sum / Math.sqrt((double) norm1 * (double) norm2));
}
+ /** vectorized cosine body */
+ private float[] cosineBody(float[] a, float[] b, int limit) {
+ int i = 0;
+ // vector loop is unrolled 4x (4 accumulators in parallel)
+ // we don't know how many the cpu can do at once, some can do 2, some 4
+ FloatVector sum1 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector sum2 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector sum3 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector sum4 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm1_1 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm1_2 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm1_3 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm1_4 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm2_1 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm2_2 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm2_3 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector norm2_4 = FloatVector.zero(FLOAT_SPECIES);
+ int unrolledLimit = limit - 3 * FLOAT_SPECIES.length();
+ for (; i < unrolledLimit; i += 4 * FLOAT_SPECIES.length()) {
+ // one
+ FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
+ FloatVector vb = FloatVector.fromArray(FLOAT_SPECIES, b, i);
+ sum1 = sum1.add(va.mul(vb));
+ norm1_1 = norm1_1.add(va.mul(va));
+ norm2_1 = norm2_1.add(vb.mul(vb));
+
+ // two
+ FloatVector vc = FloatVector.fromArray(FLOAT_SPECIES, a, i + FLOAT_SPECIES.length());
+ FloatVector vd = FloatVector.fromArray(FLOAT_SPECIES, b, i + FLOAT_SPECIES.length());
+ sum2 = sum2.add(vc.mul(vd));
+ norm1_2 = norm1_2.add(vc.mul(vc));
+ norm2_2 = norm2_2.add(vd.mul(vd));
+
+ // three
+ FloatVector ve = FloatVector.fromArray(FLOAT_SPECIES, a, i + 2 * FLOAT_SPECIES.length());
+ FloatVector vf = FloatVector.fromArray(FLOAT_SPECIES, b, i + 2 * FLOAT_SPECIES.length());
+ sum3 = sum3.add(ve.mul(vf));
+ norm1_3 = norm1_3.add(ve.mul(ve));
+ norm2_3 = norm2_3.add(vf.mul(vf));
+
+ // four
+ FloatVector vg = FloatVector.fromArray(FLOAT_SPECIES, a, i + 3 * FLOAT_SPECIES.length());
+ FloatVector vh = FloatVector.fromArray(FLOAT_SPECIES, b, i + 3 * FLOAT_SPECIES.length());
+ sum4 = sum4.add(vg.mul(vh));
+ norm1_4 = norm1_4.add(vg.mul(vg));
+ norm2_4 = norm2_4.add(vh.mul(vh));
+ }
+ // vector tail: less scalar computations for unaligned sizes, esp with big vector sizes
+ for (; i < limit; i += FLOAT_SPECIES.length()) {
+ FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
+ FloatVector vb = FloatVector.fromArray(FLOAT_SPECIES, b, i);
+ sum1 = sum1.add(va.mul(vb));
+ norm1_1 = norm1_1.add(va.mul(va));
+ norm2_1 = norm2_1.add(vb.mul(vb));
+ }
+ // reduce
+ FloatVector sumres1 = sum1.add(sum2);
+ FloatVector sumres2 = sum3.add(sum4);
+ FloatVector norm1res1 = norm1_1.add(norm1_2);
+ FloatVector norm1res2 = norm1_3.add(norm1_4);
+ FloatVector norm2res1 = norm2_1.add(norm2_2);
+ FloatVector norm2res2 = norm2_3.add(norm2_4);
+ return new float[] {
+ sumres1.add(sumres2).reduceLanes(ADD),
+ norm1res1.add(norm1res2).reduceLanes(ADD),
+ norm2res1.add(norm2res2).reduceLanes(ADD)
+ };
+ }
+
@Override
public float squareDistance(float[] a, float[] b) {
int i = 0;
float res = 0;
+
// if the array size is large (> 2x platform vector size), its worth the overhead to vectorize
- if (a.length > 2 * PREF_FLOAT_SPECIES.length()) {
- // vector loop is unrolled 4x (4 accumulators in parallel)
- FloatVector acc1 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector acc2 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector acc3 = FloatVector.zero(PREF_FLOAT_SPECIES);
- FloatVector acc4 = FloatVector.zero(PREF_FLOAT_SPECIES);
- int upperBound = PREF_FLOAT_SPECIES.loopBound(a.length - 3 * PREF_FLOAT_SPECIES.length());
- for (; i < upperBound; i += 4 * PREF_FLOAT_SPECIES.length()) {
- FloatVector va = FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i);
- FloatVector vb = FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i);
- FloatVector diff1 = va.sub(vb);
- acc1 = acc1.add(diff1.mul(diff1));
- FloatVector vc =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + PREF_FLOAT_SPECIES.length());
- FloatVector vd =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + PREF_FLOAT_SPECIES.length());
- FloatVector diff2 = vc.sub(vd);
- acc2 = acc2.add(diff2.mul(diff2));
- FloatVector ve =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + 2 * PREF_FLOAT_SPECIES.length());
- FloatVector vf =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + 2 * PREF_FLOAT_SPECIES.length());
- FloatVector diff3 = ve.sub(vf);
- acc3 = acc3.add(diff3.mul(diff3));
- FloatVector vg =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i + 3 * PREF_FLOAT_SPECIES.length());
- FloatVector vh =
- FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i + 3 * PREF_FLOAT_SPECIES.length());
- FloatVector diff4 = vg.sub(vh);
- acc4 = acc4.add(diff4.mul(diff4));
- }
- // vector tail: less scalar computations for unaligned sizes, esp with big vector sizes
- upperBound = PREF_FLOAT_SPECIES.loopBound(a.length);
- for (; i < upperBound; i += PREF_FLOAT_SPECIES.length()) {
- FloatVector va = FloatVector.fromArray(PREF_FLOAT_SPECIES, a, i);
- FloatVector vb = FloatVector.fromArray(PREF_FLOAT_SPECIES, b, i);
- FloatVector diff = va.sub(vb);
- acc1 = acc1.add(diff.mul(diff));
- }
- // reduce
- FloatVector res1 = acc1.add(acc2);
- FloatVector res2 = acc3.add(acc4);
- res += res1.add(res2).reduceLanes(VectorOperators.ADD);
+ if (a.length > 2 * FLOAT_SPECIES.length()) {
+ i += FLOAT_SPECIES.loopBound(a.length);
+ res += squareDistanceBody(a, b, i);
}
+ // scalar tail
for (; i < a.length; i++) {
float diff = a[i] - b[i];
res += diff * diff;
@@ -242,138 +253,180 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
return res;
}
+ /** vectorized square distance body */
+ private float squareDistanceBody(float[] a, float[] b, int limit) {
+ int i = 0;
+ // vector loop is unrolled 4x (4 accumulators in parallel)
+ // we don't know how many the cpu can do at once, some can do 2, some 4
+ FloatVector acc1 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector acc2 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector acc3 = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector acc4 = FloatVector.zero(FLOAT_SPECIES);
+ int unrolledLimit = limit - 3 * FLOAT_SPECIES.length();
+ for (; i < unrolledLimit; i += 4 * FLOAT_SPECIES.length()) {
+ // one
+ FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
+ FloatVector vb = FloatVector.fromArray(FLOAT_SPECIES, b, i);
+ FloatVector diff1 = va.sub(vb);
+ acc1 = acc1.add(diff1.mul(diff1));
+
+ // two
+ FloatVector vc = FloatVector.fromArray(FLOAT_SPECIES, a, i + FLOAT_SPECIES.length());
+ FloatVector vd = FloatVector.fromArray(FLOAT_SPECIES, b, i + FLOAT_SPECIES.length());
+ FloatVector diff2 = vc.sub(vd);
+ acc2 = acc2.add(diff2.mul(diff2));
+
+ // three
+ FloatVector ve = FloatVector.fromArray(FLOAT_SPECIES, a, i + 2 * FLOAT_SPECIES.length());
+ FloatVector vf = FloatVector.fromArray(FLOAT_SPECIES, b, i + 2 * FLOAT_SPECIES.length());
+ FloatVector diff3 = ve.sub(vf);
+ acc3 = acc3.add(diff3.mul(diff3));
+
+ // four
+ FloatVector vg = FloatVector.fromArray(FLOAT_SPECIES, a, i + 3 * FLOAT_SPECIES.length());
+ FloatVector vh = FloatVector.fromArray(FLOAT_SPECIES, b, i + 3 * FLOAT_SPECIES.length());
+ FloatVector diff4 = vg.sub(vh);
+ acc4 = acc4.add(diff4.mul(diff4));
+ }
+ // vector tail: less scalar computations for unaligned sizes, esp with big vector sizes
+ for (; i < limit; i += FLOAT_SPECIES.length()) {
+ FloatVector va = FloatVector.fromArray(FLOAT_SPECIES, a, i);
+ FloatVector vb = FloatVector.fromArray(FLOAT_SPECIES, b, i);
+ FloatVector diff = va.sub(vb);
+ acc1 = acc1.add(diff.mul(diff));
+ }
+ // reduce
+ FloatVector res1 = acc1.add(acc2);
+ FloatVector res2 = acc3.add(acc4);
+ return res1.add(res2).reduceLanes(ADD);
+ }
+
// Binary functions, these all follow a general pattern like this:
//
// short intermediate = a * b;
- // int accumulator = accumulator + intermediate;
+ // int accumulator = (int)accumulator + (int)intermediate;
//
// 256 or 512 bit vectors can process 64 or 128 bits at a time, respectively
// intermediate results use 128 or 256 bit vectors, respectively
// final accumulator uses 256 or 512 bit vectors, respectively
//
- // We also support 128 bit vectors, using two 128 bit accumulators.
+ // We also support 128 bit vectors, going 32 bits at a time.
// This is slower but still faster than not vectorizing at all.
@Override
public int dotProduct(byte[] a, byte[] b) {
int i = 0;
int res = 0;
+
// only vectorize if we'll at least enter the loop a single time, and we have at least 128-bit
// vectors (256-bit on intel to dodge performance landmines)
- if (a.length >= 16 && useIntegerVectors) {
+ if (a.length >= 16 && HAS_FAST_INTEGER_VECTORS) {
// compute vectorized dot product consistent with VPDPBUSD instruction
- if (INT_SPECIES_PREF_BIT_SIZE >= 256) {
- // optimized 256/512 bit implementation, processes 8/16 bytes at a time
- int upperBound = PREF_BYTE_SPECIES.loopBound(a.length);
- IntVector acc = IntVector.zero(IntVector.SPECIES_PREFERRED);
- for (; i < upperBound; i += PREF_BYTE_SPECIES.length()) {
- ByteVector va8 = ByteVector.fromArray(PREF_BYTE_SPECIES, a, i);
- ByteVector vb8 = ByteVector.fromArray(PREF_BYTE_SPECIES, b, i);
- Vector<Short> va16 = va8.convertShape(VectorOperators.B2S, PREF_SHORT_SPECIES, 0);
- Vector<Short> vb16 = vb8.convertShape(VectorOperators.B2S, PREF_SHORT_SPECIES, 0);
- Vector<Short> prod16 = va16.mul(vb16);
- Vector<Integer> prod32 =
- prod16.convertShape(VectorOperators.S2I, IntVector.SPECIES_PREFERRED, 0);
- acc = acc.add(prod32);
- }
- // reduce
- res += acc.reduceLanes(VectorOperators.ADD);
+ if (VECTOR_BITSIZE >= 512) {
+ i += BYTE_SPECIES.loopBound(a.length);
+ res += dotProductBody512(a, b, i);
+ } else if (VECTOR_BITSIZE == 256) {
+ i += BYTE_SPECIES.loopBound(a.length);
+ res += dotProductBody256(a, b, i);
} else {
- // 128-bit impl, which is tricky since we don't have SPECIES_32, it does "overlapping read"
- int upperBound = ByteVector.SPECIES_64.loopBound(a.length - ByteVector.SPECIES_64.length());
- IntVector acc = IntVector.zero(IntVector.SPECIES_128);
- // 4 bytes at a time
- for (; i < upperBound; i += ByteVector.SPECIES_64.length() >> 1) {
- // load 8 bytes
- ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
- ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
-
- // process first "half" only
- Vector<Short> va16 = va8.convert(VectorOperators.B2S, 0);
- Vector<Short> vb16 = vb8.convert(VectorOperators.B2S, 0);
- Vector<Short> prod16 = va16.mul(vb16);
-
- acc = acc.add(prod16.convertShape(VectorOperators.S2I, IntVector.SPECIES_128, 0));
- }
- // reduce
- res += acc.reduceLanes(VectorOperators.ADD);
+ // tricky: we don't have SPECIES_32, so we workaround with "overlapping read"
+ i += ByteVector.SPECIES_64.loopBound(a.length - ByteVector.SPECIES_64.length());
+ res += dotProductBody128(a, b, i);
}
}
+ // scalar tail
for (; i < a.length; i++) {
res += b[i] * a[i];
}
return res;
}
+ /** vectorized dot product body (512 bit vectors) */
+ private int dotProductBody512(byte[] a, byte[] b, int limit) {
+ IntVector acc = IntVector.zero(INT_SPECIES);
+ for (int i = 0; i < limit; i += BYTE_SPECIES.length()) {
+ ByteVector va8 = ByteVector.fromArray(BYTE_SPECIES, a, i);
+ ByteVector vb8 = ByteVector.fromArray(BYTE_SPECIES, b, i);
+
+ // 16-bit multiply: avoid AVX-512 heavy multiply on zmm
+ Vector<Short> va16 = va8.convertShape(B2S, SHORT_SPECIES, 0);
+ Vector<Short> vb16 = vb8.convertShape(B2S, SHORT_SPECIES, 0);
+ Vector<Short> prod16 = va16.mul(vb16);
+
+ // 32-bit add
+ Vector<Integer> prod32 = prod16.convertShape(S2I, INT_SPECIES, 0);
+ acc = acc.add(prod32);
+ }
+ // reduce
+ return acc.reduceLanes(ADD);
+ }
+
+ /** vectorized dot product body (256 bit vectors) */
+ private int dotProductBody256(byte[] a, byte[] b, int limit) {
+ IntVector acc = IntVector.zero(IntVector.SPECIES_256);
+ for (int i = 0; i < limit; i += ByteVector.SPECIES_64.length()) {
+ ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
+ ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
+
+ // 32-bit multiply and add into accumulator
+ Vector<Integer> va32 = va8.convertShape(B2I, IntVector.SPECIES_256, 0);
+ Vector<Integer> vb32 = vb8.convertShape(B2I, IntVector.SPECIES_256, 0);
+ acc = acc.add(va32.mul(vb32));
+ }
+ // reduce
+ return acc.reduceLanes(ADD);
+ }
+
+ /** vectorized dot product body (128 bit vectors) */
+ private int dotProductBody128(byte[] a, byte[] b, int limit) {
+ IntVector acc = IntVector.zero(IntVector.SPECIES_128);
+ // 4 bytes at a time (re-loading half the vector each time!)
+ for (int i = 0; i < limit; i += ByteVector.SPECIES_64.length() >> 1) {
+ // load 8 bytes
+ ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
+ ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
+
+ // process first "half" only: 16-bit multiply
+ Vector<Short> va16 = va8.convert(B2S, 0);
+ Vector<Short> vb16 = vb8.convert(B2S, 0);
+ Vector<Short> prod16 = va16.mul(vb16);
+
+ // 32-bit add
+ acc = acc.add(prod16.convertShape(S2I, IntVector.SPECIES_128, 0));
+ }
+ // reduce
+ return acc.reduceLanes(ADD);
+ }
+
@Override
public float cosine(byte[] a, byte[] b) {
int i = 0;
int sum = 0;
int norm1 = 0;
int norm2 = 0;
+
// only vectorize if we'll at least enter the loop a single time, and we have at least 128-bit
// vectors (256-bit on intel to dodge performance landmines)
- if (a.length >= 16 && useIntegerVectors) {
- if (INT_SPECIES_PREF_BIT_SIZE >= 256) {
- // optimized 256/512 bit implementation, processes 8/16 bytes at a time
- int upperBound = PREF_BYTE_SPECIES.loopBound(a.length);
- IntVector accSum = IntVector.zero(IntVector.SPECIES_PREFERRED);
- IntVector accNorm1 = IntVector.zero(IntVector.SPECIES_PREFERRED);
- IntVector accNorm2 = IntVector.zero(IntVector.SPECIES_PREFERRED);
- for (; i < upperBound; i += PREF_BYTE_SPECIES.length()) {
- ByteVector va8 = ByteVector.fromArray(PREF_BYTE_SPECIES, a, i);
- ByteVector vb8 = ByteVector.fromArray(PREF_BYTE_SPECIES, b, i);
- Vector<Short> va16 = va8.convertShape(VectorOperators.B2S, PREF_SHORT_SPECIES, 0);
- Vector<Short> vb16 = vb8.convertShape(VectorOperators.B2S, PREF_SHORT_SPECIES, 0);
- Vector<Short> prod16 = va16.mul(vb16);
- Vector<Short> norm1_16 = va16.mul(va16);
- Vector<Short> norm2_16 = vb16.mul(vb16);
- Vector<Integer> prod32 =
- prod16.convertShape(VectorOperators.S2I, IntVector.SPECIES_PREFERRED, 0);
- Vector<Integer> norm1_32 =
- norm1_16.convertShape(VectorOperators.S2I, IntVector.SPECIES_PREFERRED, 0);
- Vector<Integer> norm2_32 =
- norm2_16.convertShape(VectorOperators.S2I, IntVector.SPECIES_PREFERRED, 0);
- accSum = accSum.add(prod32);
- accNorm1 = accNorm1.add(norm1_32);
- accNorm2 = accNorm2.add(norm2_32);
- }
- // reduce
- sum += accSum.reduceLanes(VectorOperators.ADD);
- norm1 += accNorm1.reduceLanes(VectorOperators.ADD);
- norm2 += accNorm2.reduceLanes(VectorOperators.ADD);
+ if (a.length >= 16 && HAS_FAST_INTEGER_VECTORS) {
+ final float[] ret;
+ if (VECTOR_BITSIZE >= 512) {
+ i += BYTE_SPECIES.loopBound(a.length);
+ ret = cosineBody512(a, b, i);
+ } else if (VECTOR_BITSIZE == 256) {
+ i += BYTE_SPECIES.loopBound(a.length);
+ ret = cosineBody256(a, b, i);
} else {
- // 128-bit impl, which is tricky since we don't have SPECIES_32, it does "overlapping read"
- int upperBound = ByteVector.SPECIES_64.loopBound(a.length - ByteVector.SPECIES_64.length());
- IntVector accSum = IntVector.zero(IntVector.SPECIES_128);
- IntVector accNorm1 = IntVector.zero(IntVector.SPECIES_128);
- IntVector accNorm2 = IntVector.zero(IntVector.SPECIES_128);
- for (; i < upperBound; i += ByteVector.SPECIES_64.length() >> 1) {
- ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
- ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
-
- // process first half only
- Vector<Short> va16 = va8.convert(VectorOperators.B2S, 0);
- Vector<Short> vb16 = vb8.convert(VectorOperators.B2S, 0);
- Vector<Short> norm1_16 = va16.mul(va16);
- Vector<Short> norm2_16 = vb16.mul(vb16);
- Vector<Short> prod16 = va16.mul(vb16);
-
- // sum into accumulators
- accNorm1 =
- accNorm1.add(norm1_16.convertShape(VectorOperators.S2I, IntVector.SPECIES_128, 0));
- accNorm2 =
- accNorm2.add(norm2_16.convertShape(VectorOperators.S2I, IntVector.SPECIES_128, 0));
- accSum = accSum.add(prod16.convertShape(VectorOperators.S2I, IntVector.SPECIES_128, 0));
- }
- // reduce
- sum += accSum.reduceLanes(VectorOperators.ADD);
- norm1 += accNorm1.reduceLanes(VectorOperators.ADD);
- norm2 += accNorm2.reduceLanes(VectorOperators.ADD);
+ // tricky: we don't have SPECIES_32, so we workaround with "overlapping read"
+ i += ByteVector.SPECIES_64.loopBound(a.length - ByteVector.SPECIES_64.length());
+ ret = cosineBody128(a, b, i);
}
+ sum += ret[0];
+ norm1 += ret[1];
+ norm2 += ret[2];
}
+ // scalar tail
for (; i < a.length; i++) {
byte elem1 = a[i];
byte elem2 = b[i];
@@ -384,58 +437,153 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
return (float) (sum / Math.sqrt((double) norm1 * (double) norm2));
}
+ /** vectorized cosine body (512 bit vectors) */
+ private float[] cosineBody512(byte[] a, byte[] b, int limit) {
+ IntVector accSum = IntVector.zero(INT_SPECIES);
+ IntVector accNorm1 = IntVector.zero(INT_SPECIES);
+ IntVector accNorm2 = IntVector.zero(INT_SPECIES);
+ for (int i = 0; i < limit; i += BYTE_SPECIES.length()) {
+ ByteVector va8 = ByteVector.fromArray(BYTE_SPECIES, a, i);
+ ByteVector vb8 = ByteVector.fromArray(BYTE_SPECIES, b, i);
+
+ // 16-bit multiply: avoid AVX-512 heavy multiply on zmm
+ Vector<Short> va16 = va8.convertShape(B2S, SHORT_SPECIES, 0);
+ Vector<Short> vb16 = vb8.convertShape(B2S, SHORT_SPECIES, 0);
+ Vector<Short> norm1_16 = va16.mul(va16);
+ Vector<Short> norm2_16 = vb16.mul(vb16);
+ Vector<Short> prod16 = va16.mul(vb16);
+
+ // sum into accumulators: 32-bit add
+ Vector<Integer> norm1_32 = norm1_16.convertShape(S2I, INT_SPECIES, 0);
+ Vector<Integer> norm2_32 = norm2_16.convertShape(S2I, INT_SPECIES, 0);
+ Vector<Integer> prod32 = prod16.convertShape(S2I, INT_SPECIES, 0);
+ accNorm1 = accNorm1.add(norm1_32);
+ accNorm2 = accNorm2.add(norm2_32);
+ accSum = accSum.add(prod32);
+ }
+ // reduce
+ return new float[] {
+ accSum.reduceLanes(ADD), accNorm1.reduceLanes(ADD), accNorm2.reduceLanes(ADD)
+ };
+ }
+
+ /** vectorized cosine body (256 bit vectors) */
+ private float[] cosineBody256(byte[] a, byte[] b, int limit) {
+ IntVector accSum = IntVector.zero(IntVector.SPECIES_256);
+ IntVector accNorm1 = IntVector.zero(IntVector.SPECIES_256);
+ IntVector accNorm2 = IntVector.zero(IntVector.SPECIES_256);
+ for (int i = 0; i < limit; i += ByteVector.SPECIES_64.length()) {
+ ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
+ ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
+
+ // 16-bit multiply, and add into accumulators
+ Vector<Integer> va32 = va8.convertShape(B2I, IntVector.SPECIES_256, 0);
+ Vector<Integer> vb32 = vb8.convertShape(B2I, IntVector.SPECIES_256, 0);
+ Vector<Integer> norm1_32 = va32.mul(va32);
+ Vector<Integer> norm2_32 = vb32.mul(vb32);
+ Vector<Integer> prod32 = va32.mul(vb32);
+ accNorm1 = accNorm1.add(norm1_32);
+ accNorm2 = accNorm2.add(norm2_32);
+ accSum = accSum.add(prod32);
+ }
+ // reduce
+ return new float[] {
+ accSum.reduceLanes(ADD), accNorm1.reduceLanes(ADD), accNorm2.reduceLanes(ADD)
+ };
+ }
+
+ /** vectorized cosine body (128 bit vectors) */
+ private float[] cosineBody128(byte[] a, byte[] b, int limit) {
+ IntVector accSum = IntVector.zero(IntVector.SPECIES_128);
+ IntVector accNorm1 = IntVector.zero(IntVector.SPECIES_128);
+ IntVector accNorm2 = IntVector.zero(IntVector.SPECIES_128);
+ for (int i = 0; i < limit; i += ByteVector.SPECIES_64.length() >> 1) {
+ ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
+ ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
+
+ // process first half only: 16-bit multiply
+ Vector<Short> va16 = va8.convert(B2S, 0);
+ Vector<Short> vb16 = vb8.convert(B2S, 0);
+ Vector<Short> norm1_16 = va16.mul(va16);
+ Vector<Short> norm2_16 = vb16.mul(vb16);
+ Vector<Short> prod16 = va16.mul(vb16);
+
+ // sum into accumulators: 32-bit add
+ accNorm1 = accNorm1.add(norm1_16.convertShape(S2I, IntVector.SPECIES_128, 0));
+ accNorm2 = accNorm2.add(norm2_16.convertShape(S2I, IntVector.SPECIES_128, 0));
+ accSum = accSum.add(prod16.convertShape(S2I, IntVector.SPECIES_128, 0));
+ }
+ // reduce
+ return new float[] {
+ accSum.reduceLanes(ADD), accNorm1.reduceLanes(ADD), accNorm2.reduceLanes(ADD)
+ };
+ }
+
@Override
public int squareDistance(byte[] a, byte[] b) {
int i = 0;
int res = 0;
+
// only vectorize if we'll at least enter the loop a single time, and we have at least 128-bit
// vectors (256-bit on intel to dodge performance landmines)
- if (a.length >= 16 && useIntegerVectors) {
- if (INT_SPECIES_PREF_BIT_SIZE >= 256) {
- // optimized 256/512 bit implementation, processes 8/16 bytes at a time
- int upperBound = PREF_BYTE_SPECIES.loopBound(a.length);
- IntVector acc = IntVector.zero(IntVector.SPECIES_PREFERRED);
- for (; i < upperBound; i += PREF_BYTE_SPECIES.length()) {
- ByteVector va8 = ByteVector.fromArray(PREF_BYTE_SPECIES, a, i);
- ByteVector vb8 = ByteVector.fromArray(PREF_BYTE_SPECIES, b, i);
- Vector<Short> va16 = va8.convertShape(VectorOperators.B2S, PREF_SHORT_SPECIES, 0);
- Vector<Short> vb16 = vb8.convertShape(VectorOperators.B2S, PREF_SHORT_SPECIES, 0);
- Vector<Short> diff16 = va16.sub(vb16);
- Vector<Integer> diff32 =
- diff16.convertShape(VectorOperators.S2I, IntVector.SPECIES_PREFERRED, 0);
- acc = acc.add(diff32.mul(diff32));
- }
- // reduce
- res += acc.reduceLanes(VectorOperators.ADD);
+ if (a.length >= 16 && HAS_FAST_INTEGER_VECTORS) {
+ if (VECTOR_BITSIZE >= 256) {
+ i += BYTE_SPECIES.loopBound(a.length);
+ res += squareDistanceBody256(a, b, i);
} else {
- // 128-bit implementation, which must "split up" vectors due to widening conversions
- int upperBound = ByteVector.SPECIES_64.loopBound(a.length);
- IntVector acc1 = IntVector.zero(IntVector.SPECIES_128);
- IntVector acc2 = IntVector.zero(IntVector.SPECIES_128);
- for (; i < upperBound; i += ByteVector.SPECIES_64.length()) {
- ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
- ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
- // expand each byte vector into short vector and subtract
- Vector<Short> va16 = va8.convertShape(VectorOperators.B2S, ShortVector.SPECIES_128, 0);
- Vector<Short> vb16 = vb8.convertShape(VectorOperators.B2S, ShortVector.SPECIES_128, 0);
- Vector<Short> diff16 = va16.sub(vb16);
- // split each short vector into two int vectors, square, and add
- Vector<Integer> diff32_1 =
- diff16.convertShape(VectorOperators.S2I, IntVector.SPECIES_128, 0);
- Vector<Integer> diff32_2 =
- diff16.convertShape(VectorOperators.S2I, IntVector.SPECIES_128, 1);
- acc1 = acc1.add(diff32_1.mul(diff32_1));
- acc2 = acc2.add(diff32_2.mul(diff32_2));
- }
- // reduce
- res += acc1.add(acc2).reduceLanes(VectorOperators.ADD);
+ i += ByteVector.SPECIES_64.loopBound(a.length);
+ res += squareDistanceBody128(a, b, i);
}
}
+ // scalar tail
for (; i < a.length; i++) {
int diff = a[i] - b[i];
res += diff * diff;
}
return res;
}
+
+ /** vectorized square distance body (256+ bit vectors) */
+ private int squareDistanceBody256(byte[] a, byte[] b, int limit) {
+ IntVector acc = IntVector.zero(INT_SPECIES);
+ for (int i = 0; i < limit; i += BYTE_SPECIES.length()) {
+ ByteVector va8 = ByteVector.fromArray(BYTE_SPECIES, a, i);
+ ByteVector vb8 = ByteVector.fromArray(BYTE_SPECIES, b, i);
+
+ // 32-bit sub, multiply, and add into accumulators
+ // TODO: uses AVX-512 heavy multiply on zmm, should we just use 256-bit vectors on AVX-512?
+ Vector<Integer> va32 = va8.convertShape(B2I, INT_SPECIES, 0);
+ Vector<Integer> vb32 = vb8.convertShape(B2I, INT_SPECIES, 0);
+ Vector<Integer> diff32 = va32.sub(vb32);
+ acc = acc.add(diff32.mul(diff32));
+ }
+ // reduce
+ return acc.reduceLanes(ADD);
+ }
+
+ /** vectorized square distance body (128 bit vectors) */
+ private int squareDistanceBody128(byte[] a, byte[] b, int limit) {
+ // 128-bit implementation, which must "split up" vectors due to widening conversions
+ // it doesn't help to do the overlapping read trick, due to 32-bit multiply in the formula
+ IntVector acc1 = IntVector.zero(IntVector.SPECIES_128);
+ IntVector acc2 = IntVector.zero(IntVector.SPECIES_128);
+ for (int i = 0; i < limit; i += ByteVector.SPECIES_64.length()) {
+ ByteVector va8 = ByteVector.fromArray(ByteVector.SPECIES_64, a, i);
+ ByteVector vb8 = ByteVector.fromArray(ByteVector.SPECIES_64, b, i);
+
+ // 16-bit sub
+ Vector<Short> va16 = va8.convertShape(B2S, ShortVector.SPECIES_128, 0);
+ Vector<Short> vb16 = vb8.convertShape(B2S, ShortVector.SPECIES_128, 0);
+ Vector<Short> diff16 = va16.sub(vb16);
+
+ // 32-bit multiply and add into accumulators
+ Vector<Integer> diff32_1 = diff16.convertShape(S2I, IntVector.SPECIES_128, 0);
+ Vector<Integer> diff32_2 = diff16.convertShape(S2I, IntVector.SPECIES_128, 1);
+ acc1 = acc1.add(diff32_1.mul(diff32_1));
+ acc2 = acc2.add(diff32_2.mul(diff32_2));
+ }
+ // reduce
+ return acc1.add(acc2).reduceLanes(ADD);
+ }
}
diff --git a/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorizationProvider.java b/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorizationProvider.java
index 4292d62b7fb..b47bf802097 100644
--- a/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorizationProvider.java
+++ b/lucene/core/src/java20/org/apache/lucene/internal/vectorization/PanamaVectorizationProvider.java
@@ -21,8 +21,6 @@ import java.security.PrivilegedAction;
import java.util.Locale;
import java.util.logging.Logger;
import jdk.incubator.vector.FloatVector;
-import jdk.incubator.vector.IntVector;
-import org.apache.lucene.util.Constants;
import org.apache.lucene.util.SuppressForbidden;
/** A vectorization provider that leverages the Panama Vector API. */
@@ -30,14 +28,6 @@ final class PanamaVectorizationProvider extends VectorizationProvider {
private final VectorUtilSupport vectorUtilSupport;
- /**
- * x86 and less than 256-bit vectors.
- *
- * <p>it could be that it has only AVX1 and integer vectors are fast. it could also be that it has
- * no AVX and integer vectors are extremely slow. don't use integer vectors to avoid landmines.
- */
- private final boolean hasFastIntegerVectors;
-
// Extracted to a method to be able to apply the SuppressForbidden annotation
@SuppressWarnings("removal")
@SuppressForbidden(reason = "security manager")
@@ -45,13 +35,7 @@ final class PanamaVectorizationProvider extends VectorizationProvider {
return AccessController.doPrivileged(action);
}
- PanamaVectorizationProvider(boolean testMode) {
- final int intPreferredBitSize = IntVector.SPECIES_PREFERRED.vectorBitSize();
- if (!testMode && intPreferredBitSize < 128) {
- throw new UnsupportedOperationException(
- "Vector bit size is less than 128: " + intPreferredBitSize);
- }
-
+ PanamaVectorizationProvider() {
// hack to work around for JDK-8309727:
try {
doPrivileged(
@@ -65,20 +49,22 @@ final class PanamaVectorizationProvider extends VectorizationProvider {
"We hit initialization failure described in JDK-8309727: " + se);
}
- // check if the system is x86 and less than 256-bit vectors:
- var isAMD64withoutAVX2 = Constants.OS_ARCH.equals("amd64") && intPreferredBitSize < 256;
- this.hasFastIntegerVectors = testMode || false == isAMD64withoutAVX2;
+ if (PanamaVectorUtilSupport.VECTOR_BITSIZE < 128) {
+ throw new UnsupportedOperationException(
+ "Vector bit size is less than 128: " + PanamaVectorUtilSupport.VECTOR_BITSIZE);
+ }
- this.vectorUtilSupport = new PanamaVectorUtilSupport(hasFastIntegerVectors);
+ this.vectorUtilSupport = new PanamaVectorUtilSupport();
var log = Logger.getLogger(getClass().getName());
log.info(
String.format(
Locale.ENGLISH,
- "Java vector incubator API enabled%s; uses preferredBitSize=%d%s",
- testMode ? " (test mode)" : "",
- intPreferredBitSize,
- hasFastIntegerVectors ? "" : "; floating-point vectors only"));
+ "Java vector incubator API enabled; uses preferredBitSize=%d%s",
+ PanamaVectorUtilSupport.VECTOR_BITSIZE,
+ PanamaVectorUtilSupport.HAS_FAST_INTEGER_VECTORS
+ ? ""
+ : "; floating-point vectors only"));
}
@Override