bjoli/PersistentMap
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Added int avx dispatch in internal nodes

Browse source
This commit is contained in:
Linus Björnstam 2026-04-28 21:05:43 +02:00
parent 2207e5e87b
commit a6e8ced7f7
2 changed files with 89 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

10
PersistentMap/BTreeFunctions.cs
View file

@ -211,6 +211,16 @@ namespace PersistentMap
internal static int FindRoutingIndex<K, TStrategy>(InternalNode<K> node, K key, long keyPrefix, TStrategy strategy)
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)
{
return FallbackRoutingKeys(node.GetKeys(), key, strategy);

78
PersistentMap/KeyStrategies/IntScanner.cs
View file

@ -87,4 +87,82 @@ public static class IntScanner
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;
}
}