forked from opensearch-project/OpenSearch
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add experimental SIMD implementation of B-tree to round down dates (o…
…pensearch-project#11194) * Add experimental SIMD implementation of B-tree to round down dates Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Use system properties in favor of feature flags to remove dependency on the server module Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Removed Java 20 test sources to simplify builds Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Remove the use of forbidden APIs in unit-tests Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Migrate to the recommended usage for custom test task classpath Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Switch benchmarks module to multi-release one Signed-off-by: Andriy Redko <andriy.redko@aiven.io> * Add JMH annotation processing for JDK-20+ sources Signed-off-by: Andriy Redko <andriy.redko@aiven.io> * Make JMH annotations consistent across sources Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Improve execution of Roundable unit-tests Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Revert "Improve execution of Roundable unit-tests" This reverts commit 2e82d0a. Signed-off-by: Ketan Verma <ketan9495@gmail.com> * Add 'forced' as a possible feature flag value to simplify the execution of unit-tests Signed-off-by: Ketan Verma <ketan9495@gmail.com> --------- Signed-off-by: Ketan Verma <ketan9495@gmail.com> Signed-off-by: Andriy Redko <andriy.redko@aiven.io> Co-authored-by: Andriy Redko <andriy.redko@aiven.io> Co-authored-by: Andrew Ross <andrross@amazon.com>
- Loading branch information
Showing
9 changed files
with
370 additions
and
34 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
35 changes: 35 additions & 0 deletions
35
benchmarks/src/main/java/org/opensearch/common/round/RoundableSupplier.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.common.round; | ||
|
|
||
| import java.util.function.Supplier; | ||
|
|
||
| public class RoundableSupplier implements Supplier<Roundable> { | ||
| private final Supplier<Roundable> delegate; | ||
|
|
||
| RoundableSupplier(String type, long[] values, int size) throws ClassNotFoundException { | ||
| switch (type) { | ||
| case "binary": | ||
| delegate = () -> new BinarySearcher(values, size); | ||
| break; | ||
| case "linear": | ||
| delegate = () -> new BidirectionalLinearSearcher(values, size); | ||
| break; | ||
| case "btree": | ||
| throw new ClassNotFoundException("BtreeSearcher is not supported below JDK 20"); | ||
| default: | ||
| throw new IllegalArgumentException("invalid type: " + type); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public Roundable get() { | ||
| return delegate.get(); | ||
| } | ||
| } |
36 changes: 36 additions & 0 deletions
36
benchmarks/src/main/java20/org/opensearch/common/round/RoundableSupplier.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,36 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.common.round; | ||
|
|
||
| import java.util.function.Supplier; | ||
|
|
||
| public class RoundableSupplier implements Supplier<Roundable> { | ||
| private final Supplier<Roundable> delegate; | ||
|
|
||
| RoundableSupplier(String type, long[] values, int size) { | ||
| switch (type) { | ||
| case "binary": | ||
| delegate = () -> new BinarySearcher(values, size); | ||
| break; | ||
| case "linear": | ||
| delegate = () -> new BidirectionalLinearSearcher(values, size); | ||
| break; | ||
| case "btree": | ||
| delegate = () -> new BtreeSearcher(values, size); | ||
| break; | ||
| default: | ||
| throw new IllegalArgumentException("invalid type: " + type); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public Roundable get() { | ||
| return delegate.get(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
100 changes: 100 additions & 0 deletions
100
libs/common/src/main/java20/org/opensearch/common/round/BtreeSearcher.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,100 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.common.round; | ||
|
|
||
| import org.opensearch.common.annotation.InternalApi; | ||
|
|
||
| import jdk.incubator.vector.LongVector; | ||
| import jdk.incubator.vector.Vector; | ||
| import jdk.incubator.vector.VectorOperators; | ||
| import jdk.incubator.vector.VectorSpecies; | ||
|
|
||
| /** | ||
| * It uses vectorized B-tree search to find the round-down point. | ||
| * | ||
| * @opensearch.internal | ||
| */ | ||
| @InternalApi | ||
| class BtreeSearcher implements Roundable { | ||
| private static final VectorSpecies<Long> LONG_VECTOR_SPECIES = LongVector.SPECIES_PREFERRED; | ||
| private static final int LANES = LONG_VECTOR_SPECIES.length(); | ||
| private static final int SHIFT = log2(LANES); | ||
|
|
||
| private final long[] values; | ||
| private final long minValue; | ||
|
|
||
| BtreeSearcher(long[] values, int size) { | ||
| if (size <= 0) { | ||
| throw new IllegalArgumentException("at least one value must be present"); | ||
| } | ||
|
|
||
| int blocks = (size + LANES - 1) / LANES; // number of blocks | ||
| int length = 1 + blocks * LANES; // size of the backing array (1-indexed) | ||
|
|
||
| this.minValue = values[0]; | ||
| this.values = new long[length]; | ||
| build(values, 0, size, this.values, 1); | ||
| } | ||
|
|
||
| /** | ||
| * Builds the B-tree memory layout. | ||
| * It builds the tree recursively, following an in-order traversal. | ||
| * | ||
| * <p> | ||
| * Each block stores 'lanes' values at indices {@code i, i + 1, ..., i + lanes - 1} where {@code i} is the | ||
| * starting offset. The starting offset of the root block is 1. The branching factor is (1 + lanes) so each | ||
| * block can have these many children. Given the starting offset {@code i} of a block, the starting offset | ||
| * of its k-th child (ranging from {@code 0, 1, ..., k}) can be computed as {@code i + ((i + k) << shift)}. | ||
| * | ||
| * @param src is the sorted input array | ||
| * @param i is the index in the input array to read the value from | ||
| * @param size the number of values in the input array | ||
| * @param dst is the output array | ||
| * @param j is the index in the output array to write the value to | ||
| * @return the next index 'i' | ||
| */ | ||
| private static int build(long[] src, int i, int size, long[] dst, int j) { | ||
| if (j < dst.length) { | ||
| for (int k = 0; k < LANES; k++) { | ||
| i = build(src, i, size, dst, j + ((j + k) << SHIFT)); | ||
|
|
||
| // Fills the B-tree as a complete tree, i.e., all levels are completely filled, | ||
| // except the last level which is filled from left to right. | ||
| // The trick is to fill the destination array between indices 1...size (inclusive / 1-indexed) | ||
| // and pad the remaining array with +infinity. | ||
| dst[j + k] = (j + k <= size) ? src[i++] : Long.MAX_VALUE; | ||
| } | ||
| i = build(src, i, size, dst, j + ((j + LANES) << SHIFT)); | ||
| } | ||
| return i; | ||
| } | ||
|
|
||
| @Override | ||
| public long floor(long key) { | ||
| Vector<Long> keyVector = LongVector.broadcast(LONG_VECTOR_SPECIES, key); | ||
| int i = 1, result = 1; | ||
|
|
||
| while (i < values.length) { | ||
| Vector<Long> valuesVector = LongVector.fromArray(LONG_VECTOR_SPECIES, values, i); | ||
| int j = i + valuesVector.compare(VectorOperators.GT, keyVector).firstTrue(); | ||
| result = (j > i) ? j : result; | ||
| i += (j << SHIFT); | ||
| } | ||
|
|
||
| assert result > 1 : "key must be greater than or equal to " + minValue; | ||
| return values[result - 1]; | ||
| } | ||
|
|
||
| private static int log2(int num) { | ||
| if ((num <= 0) || ((num & (num - 1)) != 0)) { | ||
| throw new IllegalArgumentException(num + " is not a positive power of 2"); | ||
| } | ||
| return 32 - Integer.numberOfLeadingZeros(num - 1); | ||
| } | ||
| } |
57 changes: 57 additions & 0 deletions
57
libs/common/src/main/java20/org/opensearch/common/round/RoundableFactory.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,57 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| */ | ||
|
|
||
| package org.opensearch.common.round; | ||
|
|
||
| import org.opensearch.common.annotation.InternalApi; | ||
|
|
||
| import jdk.incubator.vector.LongVector; | ||
|
|
||
| /** | ||
| * Factory class to create and return the fastest implementation of {@link Roundable}. | ||
| * | ||
| * @opensearch.internal | ||
| */ | ||
| @InternalApi | ||
| public final class RoundableFactory { | ||
| /** | ||
| * The maximum limit up to which linear search is used, otherwise binary or B-tree search is used. | ||
| * This is because linear search is much faster on small arrays. | ||
| * Benchmark results: <a href="https://github.com/opensearch-project/OpenSearch/pull/9727">PR #9727</a> | ||
| */ | ||
| private static final int LINEAR_SEARCH_MAX_SIZE = 64; | ||
|
|
||
| /** | ||
| * Indicates whether the vectorized (SIMD) B-tree search implementation is to be used. | ||
| * It is true when either: | ||
| * 1. The feature flag is set to "forced", or | ||
| * 2. The platform has a minimum of 4 long vector lanes and the feature flag is set to "true". | ||
| */ | ||
| private static final boolean USE_BTREE_SEARCHER; | ||
|
|
||
| static { | ||
| String simdRoundingFeatureFlag = System.getProperty("opensearch.experimental.feature.simd.rounding.enabled"); | ||
| USE_BTREE_SEARCHER = "forced".equalsIgnoreCase(simdRoundingFeatureFlag) | ||
| || (LongVector.SPECIES_PREFERRED.length() >= 4 && "true".equalsIgnoreCase(simdRoundingFeatureFlag)); | ||
| } | ||
|
|
||
| private RoundableFactory() {} | ||
|
|
||
| /** | ||
| * Creates and returns the fastest implementation of {@link Roundable}. | ||
| */ | ||
| public static Roundable create(long[] values, int size) { | ||
| if (size <= LINEAR_SEARCH_MAX_SIZE) { | ||
| return new BidirectionalLinearSearcher(values, size); | ||
| } else if (USE_BTREE_SEARCHER) { | ||
| return new BtreeSearcher(values, size); | ||
| } else { | ||
| return new BinarySearcher(values, size); | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.