slow-J commented on code in PR #16790:
URL: https://github.com/apache/lucene/pull/16790#discussion_r4424253337
##########
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:
I used `min(max(...))` because `Math.clamp` throws when lower > upper. I
don't think anything asserts that the interval is ordered.
There is an existing
```
private static double clamp(double x, double a, double b) {
return Math.min(Math.max(x, a), b);
}
```
used in OptimizedScalarQuantizer, maybe we could use a similar helper method
here?
##########
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:
Good idea, will do this in latest commit!
##########
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:
`invokeOrDelegate` looks like it tries to invoke a C version of the
function, calling it if it exists, else falling back to the Java delegate.
There's no C version of these, so I think we should call the delegate
directly? Unless you think we should add C code here?
`invokeOrDelegate` also has a TODO saying it's slow, but I didn't look into
why.
##########
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:
Good idea, will do in next commit!
##########
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:
Will do in next commit!
##########
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:
Not my area of expertise, so if you think it's valid I can change it. Let me
know what you think.
I actually ran JMH against this before the PR was finished.
`reduceLanes(ADD)` adds the lanes in a different order before and after C2
compiles the method, so the same vector can get different quantized values
depending on JIT warm-up: 5 force merges of the same 1-bit index wrote 3
different `.veq` files.
| scalarQuantizeVector ops/ms, 2 runs | reduceLanes(ADD) | sumLanes |
|---|---:|---:|
| 1024d | 129.6 / 130.2 | 129.2 / 130.4 |
| 4096d | 30.9 / 31.2 | 30.9 / 30.8 |
##########
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:
Done in the next commit. Vectors are centered in setup with a zero centroid,
so there's no copy, and each call is checked against the setup result for that
vector.
--
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]
