-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add vectorized (SIMD) implementation of B-tree to round down dates
Signed-off-by: Ketan Verma <[email protected]>
- Loading branch information
Showing
7 changed files
with
202 additions
and
2 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
97 changes: 97 additions & 0 deletions
97
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,97 @@ | ||
/* | ||
* 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 final VectorSpecies<Long> species; | ||
private final int lanes; | ||
private final int shift; | ||
private final long[] values; | ||
|
||
BtreeSearcher(long[] values, int size, VectorSpecies<Long> species) { | ||
assert size > 0 : "at least one value must be present"; | ||
|
||
this.species = species; | ||
this.lanes = species.length(); | ||
this.shift = log2(lanes); | ||
|
||
int blocks = (size + lanes - 1) / lanes; // number of blocks | ||
int length = 1 + blocks * lanes; // size of the backing array (1-indexed) | ||
|
||
this.values = new long[length]; | ||
build(values, 0, this.values, 1); | ||
} | ||
|
||
/** | ||
* Builds the B-tree memory layout. | ||
* | ||
* <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 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 int build(long[] src, int i, long[] dst, int j) { | ||
if (j < dst.length) { | ||
for (int k = 0; k < lanes; k++) { | ||
i = build(src, i, dst, j + ((j + k) << shift)); | ||
dst[j + k] = (j + k < src.length + 1) ? src[i++] : Long.MAX_VALUE; | ||
} | ||
i = build(src, i, dst, j + ((j + lanes) << shift)); | ||
} | ||
return i; | ||
} | ||
|
||
@Override | ||
public long floor(long key) { | ||
Vector<Long> keyVector = LongVector.broadcast(species, key); | ||
int i = 1, result = 1; | ||
|
||
while (i < values.length) { | ||
Vector<Long> valuesVector = LongVector.fromArray(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 " + values[1]; | ||
return values[result - 1]; | ||
} | ||
|
||
private static int log2(int n) { | ||
assert (n > 0) && ((n & (n - 1)) == 0) : n + " is not a positive power of 2"; | ||
|
||
int result = 0; | ||
while (n > 1) { | ||
n >>>= 1; | ||
result += 1; | ||
} | ||
|
||
return result; | ||
} | ||
} |
55 changes: 55 additions & 0 deletions
55
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,55 @@ | ||
/* | ||
* 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.VectorSpecies; | ||
|
||
/** | ||
* 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 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; | ||
|
||
/** | ||
* The preferred LongVector species with the maximal bit-size supported on this platform. | ||
*/ | ||
private static final VectorSpecies<Long> LONG_VECTOR_SPECIES = LongVector.SPECIES_PREFERRED; | ||
|
||
/** | ||
* Indicates whether the vectorized (SIMD) B-tree search implementation is supported. | ||
* This is true when the platform has a minimum of 4 long vector lanes. | ||
*/ | ||
private static final boolean IS_BTREE_SEARCH_SUPPORTED = LONG_VECTOR_SPECIES.length() >= 4; | ||
|
||
private RoundableFactory() {} | ||
|
||
/** | ||
* Creates and returns the fastest implementation of {@link Roundable}. | ||
*/ | ||
public static Roundable create(long[] values, int size, boolean useSimdIfAvailable) { | ||
if (size <= LINEAR_SEARCH_MAX_SIZE) { | ||
return new BidirectionalLinearSearcher(values, size); | ||
} else if (IS_BTREE_SEARCH_SUPPORTED && useSimdIfAvailable) { | ||
return new BtreeSearcher(values, size, LONG_VECTOR_SPECIES); | ||
} else { | ||
return new BinarySearcher(values, size); | ||
} | ||
} | ||
} |
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