This is an automated email from the ASF dual-hosted git repository.
gianm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new f46ee871953 feat: simd sqrt (#19742)
f46ee871953 is described below
commit f46ee87195396a3be42303bd8c23324e2c1eba0e
Author: Clint Wylie <[email protected]>
AuthorDate: Thu Jul 23 20:25:04 2026 -0700
feat: simd sqrt (#19742)
---
.../benchmark/query/SqlExpressionBenchmark.java | 9 ++-
...VectorMathUnivariateDoubleProcessorFactory.java | 29 +++++++-
.../math/expr/vector/VectorMathProcessors.java | 2 +-
.../expr/vector/simd/SimdDoubleSqrtProcessor.java | 59 ++++++++++++++++
.../vector/simd/SimdLongToDoubleSqrtProcessor.java | 63 +++++++++++++++++
.../simd/SimdLongToDoubleUnaryProcessor.java | 82 ++++++++++++++++++++++
.../math/expr/vector/simd/SimdProcessors.java | 14 ++++
.../expr/vector/simd/SimdSupportedUnaryOp.java | 3 +-
8 files changed, 256 insertions(+), 5 deletions(-)
diff --git
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
index 43a1bb3413c..dc386debdd9 100644
---
a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
+++
b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlExpressionBenchmark.java
@@ -176,7 +176,10 @@ public class SqlExpressionBenchmark extends
SqlBaseQueryBenchmark
"SELECT SUM(ABS(long4)) FROM expressions",
"SELECT SUM(ABS(double1)) FROM expressions",
// 68: unary abs of a binary subtraction (composes SIMD sub with SIMD
abs)
- "SELECT SUM(ABS(long1 - long4)) FROM expressions"
+ "SELECT SUM(ABS(long1 - long4)) FROM expressions",
+ // 69,70: unary sqrt on double and long inputs (long input exercises the
SIMD widening path)
+ "SELECT SUM(SQRT(double1)) FROM expressions",
+ "SELECT SUM(SQRT(long1)) FROM expressions"
);
@Param({
@@ -261,7 +264,9 @@ public class SqlExpressionBenchmark extends
SqlBaseQueryBenchmark
"65",
"66",
"67",
- "68"
+ "68",
+ "69",
+ "70"
})
private String query;
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/SimpleVectorMathUnivariateDoubleProcessorFactory.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/SimpleVectorMathUnivariateDoubleProcessorFactory.java
index 593170dde3e..5d0c547de4c 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/SimpleVectorMathUnivariateDoubleProcessorFactory.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/SimpleVectorMathUnivariateDoubleProcessorFactory.java
@@ -20,32 +20,56 @@
package org.apache.druid.math.expr.vector;
import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExpressionProcessing;
import
org.apache.druid.math.expr.vector.functional.DoubleUnivariateDoubleFunction;
import
org.apache.druid.math.expr.vector.functional.DoubleUnivariateLongFunction;
+import org.apache.druid.math.expr.vector.simd.SimdProcessors;
+import org.apache.druid.math.expr.vector.simd.SimdSupportedUnaryOp;
+
+import javax.annotation.Nullable;
/**
* Make a 1 argument math processor with the following type rules
* long -> double
* double -> double
- * using simple scalar functions {@link DoubleUnivariateLongFunction} and
{@link DoubleUnivariateDoubleFunction}
+ * using simple scalar functions {@link DoubleUnivariateLongFunction} and
{@link DoubleUnivariateDoubleFunction}.
+ *
+ * If a non-null {@link SimdSupportedUnaryOp} is supplied to the constructor
and
+ * {@link ExpressionProcessing#useVectorApi()} is true, this factory will
return SIMD-specialized processors backed
+ * by the JDK incubator {@code jdk.incubator.vector} API instead of the
standard scalar implementations.
*/
public class SimpleVectorMathUnivariateDoubleProcessorFactory extends
VectorMathUnivariateDoubleProcessorFactory
{
private final DoubleUnivariateLongFunction longFunction;
private final DoubleUnivariateDoubleFunction doubleFunction;
+ @Nullable
+ private final SimdSupportedUnaryOp simdOp;
public SimpleVectorMathUnivariateDoubleProcessorFactory(
DoubleUnivariateLongFunction longFunction,
DoubleUnivariateDoubleFunction doubleFunction
)
+ {
+ this(longFunction, doubleFunction, null);
+ }
+
+ protected SimpleVectorMathUnivariateDoubleProcessorFactory(
+ DoubleUnivariateLongFunction longFunction,
+ DoubleUnivariateDoubleFunction doubleFunction,
+ @Nullable SimdSupportedUnaryOp simdOp
+ )
{
this.longFunction = longFunction;
this.doubleFunction = doubleFunction;
+ this.simdOp = simdOp;
}
@Override
public final ExprVectorProcessor<double[]>
longProcessor(Expr.VectorInputBindingInspector inspector, Expr arg)
{
+ if (simdOp != null && ExpressionProcessing.useVectorApi()) {
+ return
SimdProcessors.makeLongToDoubleUnary(arg.asVectorProcessor(inspector), simdOp,
longFunction);
+ }
return new DoubleUnivariateLongFunctionVectorProcessor(
arg.asVectorProcessor(inspector),
longFunction
@@ -55,6 +79,9 @@ public class SimpleVectorMathUnivariateDoubleProcessorFactory
extends VectorMath
@Override
public final ExprVectorProcessor<double[]>
doubleProcessor(Expr.VectorInputBindingInspector inspector, Expr arg)
{
+ if (simdOp != null && ExpressionProcessing.useVectorApi()) {
+ return SimdProcessors.makeDoubleUnary(arg.asVectorProcessor(inspector),
simdOp, doubleFunction);
+ }
return new DoubleUnivariateDoubleFunctionVectorProcessor(
arg.asVectorProcessor(inspector),
doubleFunction
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/VectorMathProcessors.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/VectorMathProcessors.java
index 17b9a629faf..86f71c662f0 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/VectorMathProcessors.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/VectorMathProcessors.java
@@ -779,7 +779,7 @@ public class VectorMathProcessors
public Sqrt()
{
- super(Math::sqrt, Math::sqrt);
+ super(Math::sqrt, Math::sqrt, SimdSupportedUnaryOp.SQRT);
}
}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdDoubleSqrtProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdDoubleSqrtProcessor.java
new file mode 100644
index 00000000000..f7512effa7a
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdDoubleSqrtProcessor.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.math.expr.vector.simd;
+
+import jdk.incubator.vector.DoubleVector;
+import jdk.incubator.vector.VectorOperators;
+import org.apache.druid.math.expr.vector.ExprVectorProcessor;
+import
org.apache.druid.math.expr.vector.functional.DoubleUnivariateDoubleFunction;
+
+import java.util.Arrays;
+
+/**
+ * SIMD specialization of {@code (double[]) -> double[]} square root. {@code
lanewise(VectorOperators.SQRT)} is
+ * emitted with the operator literally inline so the JIT statically resolves
it to the platform's double-sqrt
+ * intrinsic (e.g. {@code vsqrtpd} on x86).
+ */
+public final class SimdDoubleSqrtProcessor extends SimdDoubleUnaryProcessor
+{
+ public SimdDoubleSqrtProcessor(ExprVectorProcessor<?> input,
DoubleUnivariateDoubleFunction scalarFallback)
+ {
+ super(input, scalarFallback);
+ }
+
+ @Override
+ protected void processVector(double[] input, boolean[] inputNulls, int
currentSize)
+ {
+ final int laneCount = SPECIES.length();
+ final int upperBound = SPECIES.loopBound(currentSize);
+ int i = 0;
+ for (; i < upperBound; i += laneCount) {
+ DoubleVector.fromArray(SPECIES, input,
i).lanewise(VectorOperators.SQRT).intoArray(outValues, i);
+ }
+ for (; i < currentSize; i++) {
+ outValues[i] = scalarFallback.process(input[i]);
+ }
+ if (inputNulls == null) {
+ Arrays.fill(outNulls, 0, currentSize, false);
+ } else {
+ System.arraycopy(inputNulls, 0, outNulls, 0, currentSize);
+ }
+ }
+}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdLongToDoubleSqrtProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdLongToDoubleSqrtProcessor.java
new file mode 100644
index 00000000000..a38670d0382
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdLongToDoubleSqrtProcessor.java
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.math.expr.vector.simd;
+
+import jdk.incubator.vector.DoubleVector;
+import jdk.incubator.vector.LongVector;
+import jdk.incubator.vector.VectorOperators;
+import org.apache.druid.math.expr.vector.ExprVectorProcessor;
+import
org.apache.druid.math.expr.vector.functional.DoubleUnivariateLongFunction;
+
+import java.util.Arrays;
+
+/**
+ * SIMD specialization of {@code (long[]) -> double[]} square root. The long
input is widened lane-by-lane to a
+ * {@link DoubleVector} via {@code castShape} and then {@code
lanewise(VectorOperators.SQRT)} is emitted with the
+ * operator literally inline so the JIT statically resolves it to the
platform's double-sqrt intrinsic (e.g.
+ * {@code vsqrtpd} on x86).
+ */
+public final class SimdLongToDoubleSqrtProcessor extends
SimdLongToDoubleUnaryProcessor
+{
+ public SimdLongToDoubleSqrtProcessor(ExprVectorProcessor<?> input,
DoubleUnivariateLongFunction scalarFallback)
+ {
+ super(input, scalarFallback);
+ }
+
+ @Override
+ protected void processVector(long[] input, boolean[] inputNulls, int
currentSize)
+ {
+ final int laneCount = DOUBLE_SPECIES.length();
+ final int upperBound = DOUBLE_SPECIES.loopBound(currentSize);
+ int i = 0;
+ for (; i < upperBound; i += laneCount) {
+ final DoubleVector va =
+ (DoubleVector) LongVector.fromArray(LONG_SPECIES, input,
i).castShape(DOUBLE_SPECIES, 0);
+ va.lanewise(VectorOperators.SQRT).intoArray(outValues, i);
+ }
+ for (; i < currentSize; i++) {
+ outValues[i] = scalarFallback.process(input[i]);
+ }
+ if (inputNulls == null) {
+ Arrays.fill(outNulls, 0, currentSize, false);
+ } else {
+ System.arraycopy(inputNulls, 0, outNulls, 0, currentSize);
+ }
+ }
+}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdLongToDoubleUnaryProcessor.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdLongToDoubleUnaryProcessor.java
new file mode 100644
index 00000000000..6436b20d158
--- /dev/null
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdLongToDoubleUnaryProcessor.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.math.expr.vector.simd;
+
+import jdk.incubator.vector.DoubleVector;
+import jdk.incubator.vector.LongVector;
+import jdk.incubator.vector.VectorSpecies;
+import org.apache.druid.math.expr.Expr;
+import org.apache.druid.math.expr.ExpressionType;
+import org.apache.druid.math.expr.vector.CastToTypeVectorProcessor;
+import org.apache.druid.math.expr.vector.ExprEvalDoubleVector;
+import org.apache.druid.math.expr.vector.ExprEvalVector;
+import org.apache.druid.math.expr.vector.ExprVectorProcessor;
+import
org.apache.druid.math.expr.vector.functional.DoubleUnivariateLongFunction;
+
+import javax.annotation.Nullable;
+
+/**
+ * Abstract base for SIMD processors that compute {@code (long[]) -> double[]}
unary ops — the long input is
+ * widened lane-by-lane to a {@link DoubleVector} via {@code
castShape(DoubleVector.SPECIES_PREFERRED, 0)} in each
+ * subclass's hot loop before the op. See {@link SimdLongUnaryProcessor} for
the general design rationale.
+ */
+abstract class SimdLongToDoubleUnaryProcessor implements
ExprVectorProcessor<double[]>
+{
+ static final VectorSpecies<Long> LONG_SPECIES = LongVector.SPECIES_PREFERRED;
+ static final VectorSpecies<Double> DOUBLE_SPECIES =
DoubleVector.SPECIES_PREFERRED;
+
+ private final ExprVectorProcessor<long[]> input;
+ final DoubleUnivariateLongFunction scalarFallback;
+ final double[] outValues;
+ final boolean[] outNulls;
+
+ protected SimdLongToDoubleUnaryProcessor(
+ ExprVectorProcessor<?> input,
+ DoubleUnivariateLongFunction scalarFallback
+ )
+ {
+ this.input = CastToTypeVectorProcessor.cast(input, ExpressionType.LONG);
+ this.scalarFallback = scalarFallback;
+ this.outValues = new double[this.input.maxVectorSize()];
+ this.outNulls = new boolean[this.input.maxVectorSize()];
+ }
+
+ @Override
+ public final ExprEvalVector<double[]> evalVector(Expr.VectorInputBinding
bindings)
+ {
+ final ExprEvalVector<long[]> lhs = input.evalVector(bindings);
+ processVector(lhs.values(), lhs.getNullVector(),
bindings.getCurrentVectorSize());
+ return new ExprEvalDoubleVector(outValues, outNulls);
+ }
+
+ protected abstract void processVector(long[] input, @Nullable boolean[]
inputNulls, int currentSize);
+
+ @Override
+ public final ExpressionType getOutputType()
+ {
+ return ExpressionType.DOUBLE;
+ }
+
+ @Override
+ public final int maxVectorSize()
+ {
+ return outValues.length;
+ }
+}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdProcessors.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdProcessors.java
index c7ce4878ee8..a6f820eb7ce 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdProcessors.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdProcessors.java
@@ -25,6 +25,7 @@ import
org.apache.druid.math.expr.vector.functional.DoubleBivariateDoubleLongFun
import
org.apache.druid.math.expr.vector.functional.DoubleBivariateDoublesFunction;
import
org.apache.druid.math.expr.vector.functional.DoubleBivariateLongDoubleFunction;
import
org.apache.druid.math.expr.vector.functional.DoubleUnivariateDoubleFunction;
+import
org.apache.druid.math.expr.vector.functional.DoubleUnivariateLongFunction;
import org.apache.druid.math.expr.vector.functional.LongBivariateLongsFunction;
import org.apache.druid.math.expr.vector.functional.LongUnivariateLongFunction;
@@ -124,7 +125,20 @@ public final class SimdProcessors
return switch (op) {
case NEG -> new SimdDoubleNegProcessor(input, scalarFallback);
case ABS -> new SimdDoubleAbsProcessor(input, scalarFallback);
+ case SQRT -> new SimdDoubleSqrtProcessor(input, scalarFallback);
default -> throw DruidException.defensive("Unsupported SIMD unary
op[%s]", op);
};
}
+
+ public static ExprVectorProcessor<double[]> makeLongToDoubleUnary(
+ ExprVectorProcessor<?> input,
+ SimdSupportedUnaryOp op,
+ DoubleUnivariateLongFunction scalarFallback
+ )
+ {
+ return switch (op) {
+ case SQRT -> new SimdLongToDoubleSqrtProcessor(input, scalarFallback);
+ default -> throw DruidException.defensive("Unsupported SIMD unary op[%s]
for long->double", op);
+ };
+ }
}
diff --git
a/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdSupportedUnaryOp.java
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdSupportedUnaryOp.java
index 260fd24d2ac..f7f7e4bc498 100644
---
a/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdSupportedUnaryOp.java
+++
b/processing/src/main/java/org/apache/druid/math/expr/vector/simd/SimdSupportedUnaryOp.java
@@ -31,5 +31,6 @@ package org.apache.druid.math.expr.vector.simd;
public enum SimdSupportedUnaryOp
{
NEG,
- ABS
+ ABS,
+ SQRT
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
