slow-J commented on code in PR #16790:
URL: https://github.com/apache/lucene/pull/16790#discussion_r4426153892
##########
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++) {
Review Comment:
We can look this as a follow up if thats ok.
##########
lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorUtilBenchmark.java:
##########
@@ -118,6 +126,21 @@ public void init() {
floatsB[i] = random.nextFloat();
shortsB[i] = Float.floatToFloat16(floatsB[i]);
}
+ // Rotate through 64 vectors, as the number of descent steps depends on
the vector. They are
+ // centered here and the centroid is zero, so scalarQuantize leaves them
as they are and every
+ // call can be checked against the first.
+ quantizeVectors = new float[64][size];
+ expectedQuantizeResults = new
OptimizedScalarQuantizer.QuantizationResult[64];
+ quantizeCentroid = new float[size];
+ quantizeDest = new byte[size];
+ quantizer = new
OptimizedScalarQuantizer(VectorSimilarityFunction.DOT_PRODUCT);
+ for (int j = 0; j < 64; j++) {
+ for (int i = 0; i < size; ++i) {
+ quantizeVectors[j][i] = random.nextFloat() - floatsB[i];
+ }
+ expectedQuantizeResults[j] =
+ quantizer.scalarQuantize(quantizeVectors[j], quantizeDest, (byte) 4,
quantizeCentroid);
Review Comment:
Makes sense, I'll remove.
##########
lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorUtilBenchmark.java:
##########
@@ -118,6 +126,21 @@ public void init() {
floatsB[i] = random.nextFloat();
shortsB[i] = Float.floatToFloat16(floatsB[i]);
}
+ // Rotate through 64 vectors, as the number of descent steps depends on
the vector. They are
+ // centered here and the centroid is zero, so scalarQuantize leaves them
as they are and every
+ // call can be checked against the first.
+ quantizeVectors = new float[64][size];
+ expectedQuantizeResults = new
OptimizedScalarQuantizer.QuantizationResult[64];
+ quantizeCentroid = new float[size];
Review Comment:
Currently, setup subtracts a random centroid from each vector, and
scalarQuantize subtracts the centroid in place, so a random one here would
change the vectors on every call and we'd need to copy. Should we still do it?
##########
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:
Sure, I'll add the three to `dotProduct.h` and call them through
`invokeOrDelegate`.
##########
lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java:
##########
@@ -459,6 +459,87 @@ private float quantizeFloat(float v, byte[] dest, int
destIndex) {
}
}
+ @Override
+ public void osqCenter(float[] vector, float[] centroid, float[] stats) {
+ osqCenterImpl(vector, centroid, 0, stats);
+ }
+
+ // The Panama version calls this for its tail, with start where its vector
loop stopped.
+ static void osqCenterImpl(float[] vector, float[] centroid, int start,
float[] stats) {
+ float dot = 0;
+ float min = Float.MAX_VALUE;
+ float max = -Float.MAX_VALUE;
+ float sumSq = 0;
+ float sum = 0;
+ for (int i = start; 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, 0, lower, upper, points, stats);
+ }
+
+ // The Panama version calls this for its tail, with start where its vector
loop stopped.
+ static void osqGridStatsImpl(
+ float[] vector, int start, 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;
Review Comment:
Good idea, I'll do that in the 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) {
Review Comment:
Thanks, I'll make them experimental in the next commit.
##########
lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java:
##########
@@ -459,6 +459,87 @@ private float quantizeFloat(float v, byte[] dest, int
destIndex) {
}
}
+ @Override
+ public void osqCenter(float[] vector, float[] centroid, float[] stats) {
+ osqCenterImpl(vector, centroid, 0, stats);
+ }
+
+ // The Panama version calls this for its tail, with start where its vector
loop stopped.
+ static void osqCenterImpl(float[] vector, float[] centroid, int start,
float[] stats) {
+ float dot = 0;
+ float min = Float.MAX_VALUE;
+ float max = -Float.MAX_VALUE;
+ float sumSq = 0;
+ float sum = 0;
+ for (int i = start; 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, 0, lower, upper, points, stats);
+ }
+
+ // The Panama version calls this for its tail, with start where its vector
loop stopped.
+ static void osqGridStatsImpl(
+ float[] vector, int start, 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;
Review Comment:
Very true, I'll make the scalar sums float.
--
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]
