kaivalnp commented on code in PR #16790:
URL: https://github.com/apache/lucene/pull/16790#discussion_r4415952295
##########
lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java:
##########
@@ -459,6 +459,80 @@ private float quantizeFloat(float v, byte[] dest, int
destIndex) {
}
}
+ @Override
+ public void osqCenter(float[] vector, float[] centroid, float[] stats) {
+ float dot = 0;
+ float min = Float.MAX_VALUE;
+ float max = -Float.MAX_VALUE;
+ float sumSq = 0;
+ float sum = 0;
+ for (int i = 0; i < vector.length; i++) {
+ dot += vector[i] * centroid[i];
+ float x = vector[i] - centroid[i];
+ vector[i] = x;
+ min = Math.min(min, x);
+ max = Math.max(max, x);
+ sumSq += x * x;
+ sum += x;
+ }
+ stats[0] = dot;
+ stats[1] = min;
+ stats[2] = max;
+ stats[3] = sumSq;
+ stats[4] = sum;
+ }
+
+ @Override
+ public void osqGridStats(float[] vector, float lower, float upper, int
points, double[] stats) {
+ osqGridStatsImpl(vector, lower, upper, points, stats);
+ }
+
+ static void osqGridStatsImpl(
+ float[] vector, float lower, float upper, int points, double[] stats) {
+ float stepInv = (points - 1f) / (upper - lower);
+ float step = (upper - lower) / (points - 1f);
+ long sumK = 0;
+ long sumKK = 0;
+ double sumXK = 0;
+ double sumXD = 0;
+ double sumDD = 0;
+ double sumX = 0;
+ for (float x : vector) {
+ // Adds 0.5 and truncates, as Math.round is costly to match in the
Panama version.
+ int k = (int) ((Math.min(Math.max(x, lower), upper) - lower) * stepInv +
0.5f);
+ float d = x - (lower + k * step);
+ sumK += k;
+ sumKK += (long) k * k;
+ sumXK += x * k;
+ sumXD += x * d;
+ sumDD += d * d;
+ sumX += x;
+ }
+ stats[0] = sumK;
+ stats[1] = sumKK;
+ stats[2] = sumXK;
+ stats[3] = sumXD;
+ stats[4] = sumDD;
+ stats[5] = sumX;
+ }
+
+ @Override
+ public int osqAssign(float[] vector, float lower, float upper, float step,
byte[] dest) {
+ return osqAssignImpl(vector, 0, lower, upper, step, dest);
+ }
+
+ // The Panama version calls this for its tail, with start where its vector
loop stopped.
+ static int osqAssignImpl(
+ float[] vector, int start, float lower, float upper, float step, byte[]
dest) {
+ int sum = 0;
+ for (int i = start; i < vector.length; i++) {
+ int k = Math.round((Math.min(Math.max(vector[i], lower), upper) - lower)
/ step);
Review Comment:
Can we use `Math.clamp` instead of `Math.min(Math.max(...`?
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java:
##########
@@ -1584,6 +1584,151 @@ public float recalculateScalarQuantizationOffset(
return correction;
}
+ // reduceLanes(ADD) leaves the order of the additions open, and it adds in a
different order
+ // before C2 compiles it than after, so the codes would depend on JIT state.
+ // Adding the lanes in order keeps them the same.
+ private static float sumLanes(FloatVector v) {
+ float sum = 0;
+ for (int i = 0; i < v.length(); i++) {
+ sum += v.lane(i);
+ }
+ return sum;
+ }
Review Comment:
IMO we can accept the arbitrary order of floating point additions here, and
use the more performant `reduceLanes(ADD)`.
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java:
##########
@@ -1584,6 +1584,151 @@ public float recalculateScalarQuantizationOffset(
return correction;
}
+ // reduceLanes(ADD) leaves the order of the additions open, and it adds in a
different order
+ // before C2 compiles it than after, so the codes would depend on JIT state.
+ // Adding the lanes in order keeps them the same.
+ private static float sumLanes(FloatVector v) {
+ float sum = 0;
+ for (int i = 0; i < v.length(); i++) {
+ sum += v.lane(i);
+ }
+ return sum;
+ }
+
+ @Override
+ public void osqCenter(float[] vector, float[] centroid, float[] stats) {
+ FloatVector accDot = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accMin = FloatVector.broadcast(FLOAT_SPECIES, Float.MAX_VALUE);
+ FloatVector accMax = FloatVector.broadcast(FLOAT_SPECIES,
-Float.MAX_VALUE);
+ FloatVector accSq = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accSum = FloatVector.zero(FLOAT_SPECIES);
+ final int bound = FLOAT_SPECIES.loopBound(vector.length);
+ int i = 0;
+ for (; i < bound; i += FLOAT_SPECIES.length()) {
+ FloatVector v = FloatVector.fromArray(FLOAT_SPECIES, vector, i);
+ FloatVector c = FloatVector.fromArray(FLOAT_SPECIES, centroid, i);
+ accDot = fma(v, c, accDot);
+ FloatVector x = v.sub(c);
+ x.intoArray(vector, i);
+ accMin = accMin.min(x);
+ accMax = accMax.max(x);
+ accSq = fma(x, x, accSq);
+ accSum = accSum.add(x);
+ }
+ // scalar tail
+ float dot = 0;
+ float min = accMin.reduceLanes(VectorOperators.MIN);
+ float max = accMax.reduceLanes(VectorOperators.MAX);
+ float sumSq = 0;
+ float sum = 0;
+ for (; i < vector.length; i++) {
+ dot += vector[i] * centroid[i];
+ float x = vector[i] - centroid[i];
+ vector[i] = x;
+ min = Math.min(min, x);
+ max = Math.max(max, x);
+ sumSq += x * x;
+ sum += x;
+ }
+ stats[0] = dot + sumLanes(accDot);
+ stats[1] = min;
+ stats[2] = max;
+ stats[3] = sumSq + sumLanes(accSq);
+ stats[4] = sum + sumLanes(accSum);
+ }
+
+ // k * k is at most 255^2, so an int lane can sum about 33K of them before
it overflows. Below
+ // 2^16 dims, no lane gets that many.
+ private static final int OSQ_GRID_STATS_MAX_DIMS = 1 << 16;
+
+ @Override
+ public void osqGridStats(float[] vector, float lower, float upper, int
points, double[] stats) {
+ if (vector.length >= OSQ_GRID_STATS_MAX_DIMS) {
+ DefaultVectorUtilSupport.osqGridStatsImpl(vector, lower, upper, points,
stats);
+ return;
+ }
+ float stepInv = (points - 1f) / (upper - lower);
+ float step = (upper - lower) / (points - 1f);
+ IntVector accK = IntVector.zero(INT_SPECIES);
+ IntVector accKK = IntVector.zero(INT_SPECIES);
+ FloatVector accXK = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accXD = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accDD = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accX = FloatVector.zero(FLOAT_SPECIES);
+ final int bound = FLOAT_SPECIES.loopBound(vector.length);
+ int i = 0;
+ for (; i < bound; i += FLOAT_SPECIES.length()) {
+ FloatVector x = FloatVector.fromArray(FLOAT_SPECIES, vector, i);
+ // Not fma, so k matches the scalar version.
+ FloatVector t =
x.max(lower).min(upper).sub(lower).mul(stepInv).add(0.5f);
+ IntVector k = (IntVector) t.convert(VectorOperators.F2I, 0);
+ FloatVector kf = (FloatVector) k.convert(VectorOperators.I2F, 0);
+ FloatVector d = x.sub(kf.mul(step).add(lower));
+ accK = accK.add(k);
+ accKK = accKK.add(k.mul(k));
+ accXK = fma(x, kf, accXK);
+ accXD = fma(x, d, accXD);
+ accDD = fma(d, d, accDD);
+ accX = accX.add(x);
+ }
+ long sumK = accK.reduceLanesToLong(ADD);
+ long sumKK = accKK.reduceLanesToLong(ADD);
+ double sumXK = sumLanes(accXK);
+ double sumXD = sumLanes(accXD);
+ double sumDD = sumLanes(accDD);
+ double sumX = sumLanes(accX);
+ // scalar tail
+ for (; i < vector.length; i++) {
+ float x = vector[i];
+ int k = (int) ((Math.min(Math.max(x, lower), upper) - lower) * stepInv +
0.5f);
+ float d = x - (lower + k * step);
+ sumK += k;
+ sumKK += (long) k * k;
+ sumXK += x * k;
+ sumXD += x * d;
+ sumDD += d * d;
+ sumX += x;
+ }
+ stats[0] = sumK;
+ stats[1] = sumKK;
+ stats[2] = sumXK;
+ stats[3] = sumXD;
+ stats[4] = sumDD;
+ stats[5] = sumX;
+ }
+
+ @Override
+ public int osqAssign(float[] vector, float lower, float upper, float step,
byte[] dest) {
+ // The vector loop converts levels with F2I, so it needs every level in
[0, 2^31).
+ float maxLevel = (upper - lower) / step;
+ if ((maxLevel > 0 && maxLevel < 0x1p31f) == false) {
+ return DefaultVectorUtilSupport.osqAssignImpl(vector, 0, lower, upper,
step, dest);
+ }
+ IntVector sums = IntVector.zero(INT_SPECIES);
+ final int lanes = FLOAT_SPECIES.length();
+ // Four int vectors of levels fill one byte vector.
+ final int bound = BYTE_SPECIES_FULL.loopBound(vector.length);
+ int i = 0;
+ for (; i < bound; i += BYTE_SPECIES_FULL.length()) {
+ ByteVector out = ByteVector.zero(BYTE_SPECIES_FULL);
+ for (int part = 0; part < 4; part++) {
+ FloatVector x = FloatVector.fromArray(FLOAT_SPECIES, vector, i + part
* lanes);
+ FloatVector t = x.max(lower).min(upper).sub(lower).div(step);
+ // Same as Math.round(t): truncate, then add 1 when twice the fraction
truncates to 1. NaN
+ // converts to 0, as in Math.round.
+ IntVector k = (IntVector) t.convert(VectorOperators.F2I, 0);
+ FloatVector frac = t.sub((FloatVector) k.convert(VectorOperators.I2F,
0));
+ k = k.add((IntVector) frac.add(frac).convert(VectorOperators.F2I, 0));
Review Comment:
If we're looking for `k = Math.round(t)` where `t` is non-negative, I wonder
if we can simplify this to `k = t.add(0.5f).convert(VectorOperators.F2I, 0)`?
I see `osqGridStats` already does this..
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/NativeVectorUtilSupport.java:
##########
@@ -571,6 +571,21 @@ public float recalculateScalarQuantizationOffset(
vector.length);
}
+ @Override
+ public void osqCenter(float[] vector, float[] centroid, float[] stats) {
+ delegateVectorUtilSupport.osqCenter(vector, centroid, stats);
Review Comment:
Should we have the same `invokeOrDelegate` pattern here?
##########
lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java:
##########
@@ -192,25 +176,15 @@ public QuantizationResult scalarQuantize(
assert bits > 0 && bits <= 8;
float[] intervalScratch = new float[2];
int points = 1 << bits;
- double vecMean = 0;
- double vecVar = 0;
- float norm2 = 0;
- float centroidDot = 0;
- float min = Float.MAX_VALUE;
- float max = -Float.MAX_VALUE;
- for (int i = 0; i < vector.length; ++i) {
- if (similarityFunction != EUCLIDEAN) {
- centroidDot += vector[i] * centroid[i];
- }
- vector[i] = vector[i] - centroid[i];
- min = Math.min(min, vector[i]);
- max = Math.max(max, vector[i]);
- norm2 += (vector[i] * vector[i]);
- double delta = vector[i] - vecMean;
- vecMean += delta / (i + 1);
- vecVar += delta * (vector[i] - vecMean);
- }
- vecVar /= vector.length;
+ float[] centerStats = new float[5];
+ VectorUtil.osqCenter(vector, centroid, centerStats);
+ float centroidDot = centerStats[0];
+ float min = centerStats[1];
+ float max = centerStats[2];
+ float norm2 = centerStats[3];
+ double vecMean = centerStats[4] / (double) vector.length;
+ // Rounding can push the mean of squares minus the squared mean just below
zero.
+ double vecVar = Math.max(0, norm2 / (double) vector.length - vecMean *
vecMean);
Review Comment:
Earlier, `vecVar` was computed using Welford's algorithm -- I wonder if this
simpler version can degrade for some datasets?
I don't have a concrete example here, but should we measure a SIMD version
of the original algorithm too? (and keep using that if not too costly)
##########
lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java:
##########
@@ -1584,6 +1584,151 @@ public float recalculateScalarQuantizationOffset(
return correction;
}
+ // reduceLanes(ADD) leaves the order of the additions open, and it adds in a
different order
+ // before C2 compiles it than after, so the codes would depend on JIT state.
+ // Adding the lanes in order keeps them the same.
+ private static float sumLanes(FloatVector v) {
+ float sum = 0;
+ for (int i = 0; i < v.length(); i++) {
+ sum += v.lane(i);
+ }
+ return sum;
+ }
+
+ @Override
+ public void osqCenter(float[] vector, float[] centroid, float[] stats) {
+ FloatVector accDot = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accMin = FloatVector.broadcast(FLOAT_SPECIES, Float.MAX_VALUE);
+ FloatVector accMax = FloatVector.broadcast(FLOAT_SPECIES,
-Float.MAX_VALUE);
+ FloatVector accSq = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accSum = FloatVector.zero(FLOAT_SPECIES);
+ final int bound = FLOAT_SPECIES.loopBound(vector.length);
+ int i = 0;
+ for (; i < bound; i += FLOAT_SPECIES.length()) {
+ FloatVector v = FloatVector.fromArray(FLOAT_SPECIES, vector, i);
+ FloatVector c = FloatVector.fromArray(FLOAT_SPECIES, centroid, i);
+ accDot = fma(v, c, accDot);
+ FloatVector x = v.sub(c);
+ x.intoArray(vector, i);
+ accMin = accMin.min(x);
+ accMax = accMax.max(x);
+ accSq = fma(x, x, accSq);
+ accSum = accSum.add(x);
+ }
+ // scalar tail
+ float dot = 0;
+ float min = accMin.reduceLanes(VectorOperators.MIN);
+ float max = accMax.reduceLanes(VectorOperators.MAX);
+ float sumSq = 0;
+ float sum = 0;
+ for (; i < vector.length; i++) {
+ dot += vector[i] * centroid[i];
+ float x = vector[i] - centroid[i];
+ vector[i] = x;
+ min = Math.min(min, x);
+ max = Math.max(max, x);
+ sumSq += x * x;
+ sum += x;
+ }
+ stats[0] = dot + sumLanes(accDot);
+ stats[1] = min;
+ stats[2] = max;
+ stats[3] = sumSq + sumLanes(accSq);
+ stats[4] = sum + sumLanes(accSum);
+ }
+
+ // k * k is at most 255^2, so an int lane can sum about 33K of them before
it overflows. Below
+ // 2^16 dims, no lane gets that many.
+ private static final int OSQ_GRID_STATS_MAX_DIMS = 1 << 16;
+
+ @Override
+ public void osqGridStats(float[] vector, float lower, float upper, int
points, double[] stats) {
+ if (vector.length >= OSQ_GRID_STATS_MAX_DIMS) {
+ DefaultVectorUtilSupport.osqGridStatsImpl(vector, lower, upper, points,
stats);
+ return;
+ }
+ float stepInv = (points - 1f) / (upper - lower);
+ float step = (upper - lower) / (points - 1f);
+ IntVector accK = IntVector.zero(INT_SPECIES);
+ IntVector accKK = IntVector.zero(INT_SPECIES);
+ FloatVector accXK = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accXD = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accDD = FloatVector.zero(FLOAT_SPECIES);
+ FloatVector accX = FloatVector.zero(FLOAT_SPECIES);
+ final int bound = FLOAT_SPECIES.loopBound(vector.length);
+ int i = 0;
+ for (; i < bound; i += FLOAT_SPECIES.length()) {
+ FloatVector x = FloatVector.fromArray(FLOAT_SPECIES, vector, i);
+ // Not fma, so k matches the scalar version.
+ FloatVector t =
x.max(lower).min(upper).sub(lower).mul(stepInv).add(0.5f);
+ IntVector k = (IntVector) t.convert(VectorOperators.F2I, 0);
+ FloatVector kf = (FloatVector) k.convert(VectorOperators.I2F, 0);
+ FloatVector d = x.sub(kf.mul(step).add(lower));
+ accK = accK.add(k);
+ accKK = accKK.add(k.mul(k));
+ accXK = fma(x, kf, accXK);
+ accXD = fma(x, d, accXD);
+ accDD = fma(d, d, accDD);
+ accX = accX.add(x);
+ }
+ long sumK = accK.reduceLanesToLong(ADD);
+ long sumKK = accKK.reduceLanesToLong(ADD);
+ double sumXK = sumLanes(accXK);
+ double sumXD = sumLanes(accXD);
+ double sumDD = sumLanes(accDD);
+ double sumX = sumLanes(accX);
+ // scalar tail
+ for (; i < vector.length; i++) {
+ float x = vector[i];
+ int k = (int) ((Math.min(Math.max(x, lower), upper) - lower) * stepInv +
0.5f);
+ float d = x - (lower + k * step);
+ sumK += k;
+ sumKK += (long) k * k;
+ sumXK += x * k;
+ sumXD += x * d;
+ sumDD += d * d;
+ sumX += x;
+ }
Review Comment:
This appears to be a duplicate of the non-SIMD version, should we refactor
into a reusable static method? (same for `osqCenter`)
##########
lucene/core/src/java/org/apache/lucene/util/VectorUtil.java:
##########
@@ -598,6 +598,52 @@ public static float recalculateOffset(
vector, oldAlpha, oldMinQuantile, scale, alpha, minQuantile,
maxQuantile);
}
+ /**
+ * Subtracts {@code centroid} from {@code vector} in place, for the
optimized scalar quantizer,
+ * and writes into {@code stats} the dot product of the original vector with
{@code centroid},
+ * then the min, max, sum of squares and sum of the result.
+ */
+ public static void osqCenter(float[] vector, float[] centroid, float[]
stats) {
+ if (vector.length != centroid.length) {
+ throw new IllegalArgumentException(
+ "vector dimensions differ: " + vector.length + "!=" +
centroid.length);
+ }
+ if (stats.length < 5) {
Review Comment:
Can we make these checks stricter? (i.e. `== 5`)
##########
lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorUtilBenchmark.java:
##########
@@ -304,6 +321,20 @@ public byte[] binaryHalfByteUnpackVector() {
return halfBytesUnpackDest;
}
+ // scalarQuantize centers its input in place, so each call copies the next
vector first.
+ @Benchmark
+ public OptimizedScalarQuantizer.QuantizationResult scalarQuantizeScalar() {
+ System.arraycopy(quantizeVectors[quantizeNext++ & 63], 0, quantizeScratch,
0, size);
+ return quantizer.scalarQuantize(quantizeScratch, quantizeDest, (byte) 4,
floatsB);
Review Comment:
Can we add some kind of check to ensure that the value produced is correct?
I realize this is costlier to do than other functions like
`binaryHalfByteDotProductSinglePackedScalar`, but might be helpful to catch
errors early while tweaking the function.
Also, should the `arraycopy` be done outside the measurement?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
