Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor BITPOS Implementation #1016

Merged
merged 10 commits into from
Feb 14, 2025
Merged

Refactor BITPOS Implementation #1016

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
a6a4755
wip
vazois Feb 9, 2025
93f8031
cleanup bitpos tests
vazois Feb 10, 2025
c4d86b3
Fix bitpos bit search
vazois Feb 11, 2025
db835bc
simplify bitpos byte search
vazois Feb 12, 2025
ea0579c
cleanup unit tests for bitpos
vazois Feb 12, 2025
5578db4
fix masked invalidPayload matching
vazois Feb 12, 2025
75368aa
add more tests
vazois Feb 13, 2025
3666295
fix boundary conditions for bit search and add tests
vazois Feb 13, 2025
bb7ad0f
change to portable intrinsics
vazois Feb 13, 2025
62ece05
add more tests for bitpos
vazois Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
267 changes: 128 additions & 139 deletions libs/server/Resp/Bitmap/BitmapManagerBitPos.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,185 +1,174 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT license.

using System.Diagnostics;
using System.Runtime.Intrinsics.X86;
using System.Buffers.Binary;
using System.Numerics;

namespace Garnet.server
{
public unsafe partial class BitmapManager
{
/// <summary>
/// Find pos of bit set/clear for given bit offsets within a single byte.
/// Main driver for BITPOS command
/// </summary>
/// <param name="value">Byte value to search within.</param>
/// <param name="bSetVal">Bit value to search for (0|1).</param>
/// <param name="startBitOffset">Start most significant bit offset in byte value.</param>
/// <param name="endBitOffset">End most significant bit offset in bitmap.</param>
/// <returns></returns>
private static long BitPosIndexBitSingleByteSearch(byte value, byte bSetVal, int startBitOffset = 0, int endBitOffset = 8)
{
Debug.Assert(startBitOffset >= 0 && startBitOffset <= 8);
Debug.Assert(endBitOffset >= 0 && endBitOffset <= 8);
bool bflag = (bSetVal == 0);
long mask = bflag ? -1 : 0;

int leftBitIndex = 1 << (8 - startBitOffset);
int rightBitIndex = 1 << (8 - endBitOffset);

// Create extraction mask
long extract = leftBitIndex - rightBitIndex;

long payload = (long)(value & extract) << 56;
// Trim leading bits
payload = payload << startBitOffset;

// Transform to count leading zeros
payload = bflag ? ~payload : payload;

// Return not found
if (payload == mask) return -1;

return (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
}

/// <summary>
/// Find pos of bit set/clear for given bit offset.
/// </summary>
/// <param name="value">Pointer to start of bitmap.</param>
/// <param name="bSetVal">Bit value to search for (0|1).</param>
/// <param name="offset">Bit offset in bitmap.</param>
/// <returns></returns>
private static long BitPosIndexBitSearch(byte* value, byte bSetVal, long offset = 0)
{
bool bflag = (bSetVal == 0);
long mask = bflag ? -1 : 0;
long startByteOffset = (offset / 8);
int bitOffset = (int)(offset & 7);

long payload = (long)value[startByteOffset] << 56;
// Trim leading bits
payload = payload << bitOffset;

// Transform to count leading zeros
payload = bflag ? ~payload : payload;

// Return not found
if (payload == mask)
return -1;

return (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
}

/// <summary>
/// Main driver for bit position command.
/// </summary>
/// <param name="setVal"></param>
/// <param name="input"></param>
/// <param name="inputLen"></param>
/// <param name="startOffset"></param>
/// <param name="endOffset"></param>
/// <param name="searchFor"></param>
/// <param name="offsetType"></param>
/// <param name="value">Pointer to start of bitmap.</param>
/// <param name="valLen">Length of bitmap.</param>
/// <returns></returns>
public static long BitPosDriver(byte setVal, long startOffset, long endOffset, byte offsetType, byte* value, int valLen)
public static long BitPosDriver(byte* input, int inputLen, long startOffset, long endOffset, byte searchFor, byte offsetType)
{
if (offsetType == 0x0)
{
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, valLen) : startOffset;
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, valLen) : endOffset;
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, inputLen) : startOffset;
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, inputLen) : endOffset;

if (startOffset >= valLen) // If startOffset greater that valLen always bitpos -1
if (startOffset >= inputLen) // If startOffset greater that valLen always bitpos -1
return -1;

if (startOffset > endOffset) // If start offset beyond endOffset return 0
return -1;

endOffset = endOffset >= valLen ? valLen : endOffset;
long pos = BitPosByte(value, setVal, startOffset, endOffset);
// check if position is exceeding the last byte in acceptable range
return pos >= ((endOffset + 1) * 8) ? -1 : pos;
endOffset = endOffset >= inputLen ? inputLen : endOffset;
// BYTE search
return BitPosByteSearch(input, inputLen, startOffset, endOffset, searchFor);
}

startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, valLen * 8) : startOffset;
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, valLen * 8) : endOffset;

var startByte = (startOffset / 8);
var endByte = (endOffset / 8);
if (startByte == endByte)
else
{
// Search only inside single byte for pos
var leftBitIndex = (int)(startOffset & 7);
var rightBitIndex = (int)((endOffset + 1) & 7);
var _ipos = BitPosIndexBitSingleByteSearch(value[startByte], setVal, leftBitIndex, rightBitIndex);
return _ipos == -1 ? _ipos : startOffset + _ipos;
}
startOffset = startOffset < 0 ? ProcessNegativeOffset(startOffset, inputLen * 8) : startOffset;
endOffset = endOffset < 0 ? ProcessNegativeOffset(endOffset, inputLen * 8) : endOffset;

var startByteIndex = startOffset >> 3;
var endByteIndex = endOffset >> 3;

// Search prefix and terminate if found position of bit
var _ppos = BitPosIndexBitSearch(value, setVal, startOffset);
if (_ppos != -1) return startOffset + _ppos;
if (startByteIndex >= inputLen) // If startOffset greater that valLen always bitpos -1
return -1;

// Adjust offsets to skip first and last byte
var _startOffset = (startOffset / 8) + 1;
var _endOffset = (endOffset / 8) - 1;
var _bpos = BitPosByte(value, setVal, _startOffset, _endOffset);
if (startByteIndex > endByteIndex) // If start offset beyond endOffset return 0
return -1;

if (_bpos != -1 && _bpos < (_endOffset + 1) * 8) return _bpos;
endOffset = endByteIndex >= inputLen ? inputLen << 3 : endOffset;

// Search suffix
var _spos = BitPosIndexBitSearch(value, setVal, endOffset);
return _spos;
// BIT search
return BitPosBitSearch(input, inputLen, startOffset, endOffset, searchFor);
}
}

/// <summary>
/// Find pos of set/clear bit in a sequence of bytes.
/// Search for position of bit set in byte array using bit offset for start and end range
/// </summary>
/// <param name="value">Pointer to start of bitmap.</param>
/// <param name="bSetVal">The bit value to search for (0 for cleared bit or 1 for set bit).</param>
/// <param name="startOffset">Starting offset into bitmap.</param>
/// <param name="endOffset">End offset into bitmap.</param>
/// <param name="input"></param>
/// <param name="inputLen"></param>
/// <param name="startBitOffset"></param>
/// <param name="endBitOffset"></param>
/// <param name="searchFor"></param>
/// <returns></returns>
private static long BitPosByte(byte* value, byte bSetVal, long startOffset, long endOffset)
private static long BitPosBitSearch(byte* input, long inputLen, long startBitOffset, long endBitOffset, byte searchFor)
{
// Mask set to look for 0 or 1 depending on clear/set flag
bool bflag = (bSetVal == 0);
long mask = bflag ? -1 : 0;
long len = (endOffset - startOffset) + 1;
long remainder = len & 7;
byte* curr = value + startOffset;
byte* end = curr + (len - remainder);

// Search for first word not matching mask.
while (curr < end)
var searchBit = searchFor == 1;
var invalidPayload = (byte)(searchBit ? 0x00 : 0xff);
var currentBitOffset = (int)startBitOffset;
while (currentBitOffset <= endBitOffset)
{
long v = *(long*)(curr);
if (v != mask) break;
curr += 8;
}
var byteIndex = currentBitOffset >> 3;
var leftBitOffset = currentBitOffset & 7;
var boundary = 8 - leftBitOffset;
var rightBitOffset = currentBitOffset + boundary <= endBitOffset ? leftBitOffset + boundary : (int)(endBitOffset & 7) + 1;

// Trim byte to start and end bit index
var mask = (0xff >> leftBitOffset) ^ (0xff >> rightBitOffset);
var payload = (long)(input[byteIndex] & mask);

long pos = (((long)(curr - value)) << 3);
// Invalid only if equals the masked payload
var invalidMask = invalidPayload & mask;

long payload = 0;
// Adjust end so we can retrieve word
end = end + remainder;
// If transformed payload is invalid skip to next byte
if (payload != invalidMask)
{
payload <<= (56 + leftBitOffset);
payload = searchBit ? payload : ~payload;

// Build payload at least one byte to examine
if (curr < end) payload |= (long)curr[0] << 56;
if (curr + 1 < end) payload |= (long)curr[1] << 48;
if (curr + 2 < end) payload |= (long)curr[2] << 40;
if (curr + 3 < end) payload |= (long)curr[3] << 32;
if (curr + 4 < end) payload |= (long)curr[4] << 24;
if (curr + 5 < end) payload |= (long)curr[5] << 16;
if (curr + 6 < end) payload |= (long)curr[6] << 8;
if (curr + 7 < end) payload |= (long)curr[7];
var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
return currentBitOffset + lzcnt;
}

// Transform to count leading zeros
payload = (bSetVal == 0) ? ~payload : payload;
currentBitOffset += boundary;
}

if (payload == mask)
return pos + 0;
return -1;
}

pos += (long)Lzcnt.X64.LeadingZeroCount((ulong)payload);
/// <summary>
/// Search for position of bit set in byte array using byte offset for start and end range
/// </summary>
/// <param name="input"></param>
/// <param name="inputLen"></param>
/// <param name="startOffset"></param>
/// <param name="endOffset"></param>
/// <param name="searchFor"></param>
/// <returns></returns>
private static long BitPosByteSearch(byte* input, long inputLen, long startOffset, long endOffset, byte searchFor)
{
// Initialize variables
var searchBit = searchFor == 1;
var invalidMask8 = searchBit ? 0x00 : 0xff;
var invalidMask32 = searchBit ? 0 : -1;
var invalidMask64 = searchBit ? 0L : -1L;
var currentStartOffset = startOffset;

while (currentStartOffset <= endOffset)
{
var remainder = endOffset - currentStartOffset + 1;
if (remainder >= 8)
{
var payload = *(long*)(input + currentStartOffset);
payload = BinaryPrimitives.ReverseEndianness(payload);

// Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
if (payload != invalidMask64)
{
// Transform to count leading zeros
payload = searchBit ? payload : ~payload;
var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
return (currentStartOffset << 3) + lzcnt;
}
currentStartOffset += 8;
}
else if (remainder >= 4)
{
var payload = *(int*)(input + currentStartOffset);
payload = BinaryPrimitives.ReverseEndianness(payload);

// Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
if (payload != invalidMask32)
{
// Transform to count leading zeros
payload = searchBit ? payload : ~payload;
var lzcnt = (long)BitOperations.LeadingZeroCount((uint)payload);
return (currentStartOffset << 3) + lzcnt;
}
currentStartOffset += 4;
}
else
{
// Process only if payload is valid (i.e. not all bits are set or clear based on searchFor parameter)
if (input[currentStartOffset] != invalidMask8)
{
// Create a payload with the current byte shifted to the most significant byte position
var payload = (long)input[currentStartOffset] << 56;
// Transform to count leading zeros
payload = searchBit ? payload : ~payload;
var lzcnt = (long)BitOperations.LeadingZeroCount((ulong)payload);
return (currentStartOffset << 3) + lzcnt;
}
currentStartOffset++;
}
}

return pos;
// Return -1 if no matching bit is found
return -1;
}
}
}
10 changes: 8 additions & 2 deletions libs/server/Storage/Functions/MainStore/PrivateMethods.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,8 +184,14 @@ void CopyRespToWithInput(ref RawStringInput input, ref SpanByte value, ref SpanB
}
}

var pos = BitmapManager.BitPosDriver(bpSetVal, bpStartOffset, bpEndOffset, bpOffsetType,
value.ToPointer() + functionsState.etagState.etagSkippedStart, value.Length - functionsState.etagState.etagSkippedStart);
var pos = BitmapManager.BitPosDriver(
input: value.ToPointer() + functionsState.etagState.etagSkippedStart,
inputLen: value.Length - functionsState.etagState.etagSkippedStart,
startOffset: bpStartOffset,
endOffset: bpEndOffset,
searchFor: bpSetVal,
offsetType: bpOffsetType
);
*(long*)dst.SpanByte.ToPointer() = pos;
CopyRespNumber(pos, ref dst);
break;
Expand Down
Loading