Added int avx dispatch in internal nodes
This commit is contained in:
parent
2207e5e87b
commit
a6e8ced7f7
2 changed files with 89 additions and 1 deletions
|
|
@ -211,6 +211,16 @@ namespace PersistentMap
|
||||||
internal static int FindRoutingIndex<K, TStrategy>(InternalNode<K> node, K key, long keyPrefix, TStrategy strategy)
|
internal static int FindRoutingIndex<K, TStrategy>(InternalNode<K> node, K key, long keyPrefix, TStrategy strategy)
|
||||||
where TStrategy : IKeyStrategy<K>
|
where TStrategy : IKeyStrategy<K>
|
||||||
{
|
{
|
||||||
|
|
||||||
|
if (typeof(K) == typeof(int))
|
||||||
|
{
|
||||||
|
Span<K> keys = node.GetKeys();
|
||||||
|
ref K firstKeyRef = ref MemoryMarshal.GetReference(keys);
|
||||||
|
ref int firstIntRef = ref Unsafe.As<K, int>(ref firstKeyRef);
|
||||||
|
ReadOnlySpan<int> intKeys = MemoryMarshal.CreateReadOnlySpan(ref firstIntRef, keys.Length);
|
||||||
|
int intKey = Unsafe.As<K, int>(ref key);
|
||||||
|
return IntScanner.FindFirstGreater(intKeys, intKey);
|
||||||
|
}
|
||||||
if (!strategy.UsesPrefixes)
|
if (!strategy.UsesPrefixes)
|
||||||
{
|
{
|
||||||
return FallbackRoutingKeys(node.GetKeys(), key, strategy);
|
return FallbackRoutingKeys(node.GetKeys(), key, strategy);
|
||||||
|
|
|
||||||
|
|
@ -87,4 +87,82 @@ public static class IntScanner
|
||||||
|
|
||||||
return LinearScan(keys.Slice(i), target) + i;
|
return LinearScan(keys.Slice(i), target) + i;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
[MethodImpl(MethodImplOptions.AggressiveInlining)]
|
||||||
|
public static int FindFirstGreater(ReadOnlySpan<int> keys, int target)
|
||||||
|
{
|
||||||
|
if (!Avx2.IsSupported || keys.Length < 8)
|
||||||
|
return LinearScanGreater(keys, target);
|
||||||
|
|
||||||
|
return Avx512F.IsSupported
|
||||||
|
? ScanAvx512Greater(keys, target)
|
||||||
|
: ScanAvx2Greater(keys, target);
|
||||||
|
}
|
||||||
|
|
||||||
|
[MethodImpl(MethodImplOptions.AggressiveInlining)]
|
||||||
|
private static int LinearScanGreater(ReadOnlySpan<int> keys, int target)
|
||||||
|
{
|
||||||
|
for (var i = 0; i < keys.Length; i++)
|
||||||
|
if (keys[i] > target)
|
||||||
|
return i;
|
||||||
|
return keys.Length;
|
||||||
|
}
|
||||||
|
|
||||||
|
[MethodImpl(MethodImplOptions.AggressiveInlining)]
|
||||||
|
private static unsafe int ScanAvx2Greater(ReadOnlySpan<int> keys, int target)
|
||||||
|
{
|
||||||
|
// For > target, AVX2 CompareGreaterThan works directly without the (target - 1) offset
|
||||||
|
var vTarget = Vector256.Create(target);
|
||||||
|
var i = 0;
|
||||||
|
var len = keys.Length;
|
||||||
|
|
||||||
|
for (; i <= len - 8; i += 8)
|
||||||
|
{
|
||||||
|
fixed (int* ptr = keys)
|
||||||
|
{
|
||||||
|
var vData = Avx2.LoadVector256(ptr + i);
|
||||||
|
var vResult = Avx2.CompareGreaterThan(vData, vTarget);
|
||||||
|
|
||||||
|
var mask = (uint)Avx2.MoveMask(vResult.AsByte());
|
||||||
|
|
||||||
|
if (mask != 0)
|
||||||
|
{
|
||||||
|
return i + (BitOperations.TrailingZeroCount(mask) / 4);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return LinearScanGreater(keys.Slice(i), target) + i;
|
||||||
|
}
|
||||||
|
|
||||||
|
[MethodImpl(MethodImplOptions.AggressiveInlining)]
|
||||||
|
private static unsafe int ScanAvx512Greater(ReadOnlySpan<int> keys, int target)
|
||||||
|
{
|
||||||
|
var vTarget = Vector512.Create(target);
|
||||||
|
var i = 0;
|
||||||
|
var len = keys.Length;
|
||||||
|
|
||||||
|
for (; i <= len - 16; i += 16)
|
||||||
|
{
|
||||||
|
fixed (int* ptr = keys)
|
||||||
|
{
|
||||||
|
var vData = Avx512F.LoadVector512(ptr + i);
|
||||||
|
|
||||||
|
// Use GreaterThan instead of GreaterThanOrEqual
|
||||||
|
var mask = Vector512.GreaterThan(vData, vTarget);
|
||||||
|
|
||||||
|
if (mask != Vector512<int>.Zero)
|
||||||
|
{
|
||||||
|
uint m = (uint)mask.ExtractMostSignificantBits();
|
||||||
|
return i + BitOperations.TrailingZeroCount(m);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return LinearScanGreater(keys.Slice(i), target) + i;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue