kaivalnp commented on code in PR #16790:
URL: https://github.com/apache/lucene/pull/16790#discussion_r4424696499
##########
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:
This is a difference in precision b/w the default and SIMD version (`double`
v/s `float`); I feel that they should be consistent (apart from the documented
/ accepted variation with order of float operations).
Perhaps relaxing these to `float` is fine too, but a second opinion would be
appreciated.
##########
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:
`OptimizedScalarQuantizer` is marked
[experimental](https://github.com/apache/lucene/blob/46432a3f0db981f8d1ef5e6bdc142f0704b7c321/lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java#L40)
and this opens up its internals -- I think we should mark these functions
similarly.
I was initially concerned about exposing quantization in `VectorUtil`, but I
see there is some precedent in #13922 and #14304.
##########
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:
I wonder if for some bit widths: it becomes faster to write to the array
without accumulating 4x into a single SIMD register first. We can defer this
for later though.
##########
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:
I was hoping to add a function to the header file [like
this](https://github.com/apache/lucene/blob/208608f75307ccc19f0fb99dc8ddcce817240f66/lucene/misc/src/c/dotProduct.h#L59-L64)
and wiring up the native provider, so that implementers can override _if
needed_ (we don't need to write the entire C implementation, I see many
functions are missing from
https://github.com/apache/lucene/blob/main/lucene/misc/src/c/dotProduct.c).
##########
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:
`osqCenter` computes `sumX` internally, I wonder if it can be passed on (and
avoid recomputing here)?
Though if this is small, it might be worth keeping as-is for code clarity.
##########
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:
This sets "expected" using the same function that is being measured, so
correctness issues may not result in exceptions being thrown?
Sorry, I realize that redoing the "quantize" function is a hassle here,
perhaps it is ok to remove the check entirely (but any future changes to the
function will need to be accompanied by end-to-end checks like you did here).
##########
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:
Can this be a random vector too, just to exercise more randomness?
--
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]
