Skip to content

Optimize managed FastTree Sumup to native parity on arm64 #7670

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

Open
vladimir-aubrecht wants to merge 4 commits into
base: main
Choose a base branch
from
+276 −13
Open

Optimize managed FastTree Sumup to native parity on arm64 #7670

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0065c94
Optimize managed FastTree Sumup to native parity on arm64
Aug 10, 2026
bb33438
Address PR review: remove nested Sumup timer in SegmentIntArray.Sumup…
Aug 11, 2026
946ffe6
Add managed vs native/reference Sumup parity tests
Aug 12, 2026
16ff030
Revert "Add managed vs native/reference Sumup parity tests"
Aug 13, 2026
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
187 changes: 175 additions & 12 deletions src/Microsoft.ML.FastTree/Dataset/DenseIntArray.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,123 @@ protected static unsafe void SumupCPlusPlusDense(SumupInputData input, FeatureHi
}
}

/// <summary>
/// Managed equivalent of <see cref="SumupCPlusPlusDense"/>, used on platforms where the
/// native FastTree library is not available (e.g. arm64). This mirrors the native C_Sumup
/// loop in src/Native/FastTreeNative (Sumup.h / SumupNibbles.h) exactly, including the
/// per-document iteration and accumulation order, so histogram results are bit-identical
/// to the native implementation. Reads go through fixed pointers to avoid the per-element
/// bounds checks and interface-indexer dispatch of the generic <see cref="IntArray.Sumup"/>.
/// </summary>
protected static unsafe void SumupManagedDense(SumupInputData input, FeatureHistogram histogram,
byte* data, int numBits)
{
using (Timer.Time(TimerEvent.SumupCppDense))
{
fixed (FloatType* pSumTargetsByBin = histogram.SumTargetsByBin)
fixed (FloatType* pSampleOutputs = input.Outputs)
fixed (double* pSumWeightsByBin = histogram.SumWeightsByBin)
fixed (double* pSampleWeights = input.Weights)
fixed (int* pIndices = input.DocIndices)
fixed (int* pCountByBin = histogram.CountByBin)
{
int count = input.TotalCount;
ushort* data16 = (ushort*)data;
int* data32 = (int*)data;

// numBits is switched outside the loop (it never varies within a call) so the
// hot loop stays a tight scalar accumulation matching the native code. The
// "pIndices == null ? i : pIndices[i]" ternary is loop-invariant and free.
if (pSumWeightsByBin != null)
{
switch (numBits)
{
case 4:
for (int i = 0; i < count; i++)
{
int p = pIndices == null ? i : pIndices[i];
int featureBin = (data[p >> 1] >> ((~(p << 2)) & 4)) & 0xf;
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
pSumWeightsByBin[featureBin] += pSampleWeights[i];
++pCountByBin[featureBin];
}
break;
case 8:
for (int i = 0; i < count; i++)
{
int featureBin = data[pIndices == null ? i : pIndices[i]];
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
pSumWeightsByBin[featureBin] += pSampleWeights[i];
++pCountByBin[featureBin];
}
break;
case 16:
for (int i = 0; i < count; i++)
{
int featureBin = data16[pIndices == null ? i : pIndices[i]];
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
pSumWeightsByBin[featureBin] += pSampleWeights[i];
++pCountByBin[featureBin];
}
break;
case 32:
for (int i = 0; i < count; i++)
{
int featureBin = data32[pIndices == null ? i : pIndices[i]];
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
pSumWeightsByBin[featureBin] += pSampleWeights[i];
++pCountByBin[featureBin];
}
break;
default:
throw Contracts.Except("Unsupported bits per item {0}", numBits);
}
}
else
{
switch (numBits)
{
case 4:
for (int i = 0; i < count; i++)
{
int p = pIndices == null ? i : pIndices[i];
int featureBin = (data[p >> 1] >> ((~(p << 2)) & 4)) & 0xf;
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
++pCountByBin[featureBin];
}
break;
case 8:
for (int i = 0; i < count; i++)
{
int featureBin = data[pIndices == null ? i : pIndices[i]];
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
++pCountByBin[featureBin];
}
break;
case 16:
for (int i = 0; i < count; i++)
{
int featureBin = data16[pIndices == null ? i : pIndices[i]];
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
++pCountByBin[featureBin];
}
break;
case 32:
for (int i = 0; i < count; i++)
{
int featureBin = data32[pIndices == null ? i : pIndices[i]];
pSumTargetsByBin[featureBin] += pSampleOutputs[i];
++pCountByBin[featureBin];
}
break;
default:
throw Contracts.Except("Unsupported bits per item {0}", numBits);
}
}
}
}
}

public override IIntArrayForwardIndexer GetIndexer()
{
return this;
Expand Down Expand Up @@ -389,21 +506,21 @@ public Dense8BitIntArray(int len)
: base(len)
{
_data = new byte[len];
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public Dense8BitIntArray(byte[] buffer, ref int position)
: base(buffer.ToInt(ref position))
{
_data = buffer.ToByteArray(ref position);
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public Dense8BitIntArray(int len, IEnumerable<int> values)
: base(len)
{
_data = values.Select(i => (byte)i).ToArray(len);
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

/// <summary>
Expand Down Expand Up @@ -457,6 +574,17 @@ private void SumupNative(SumupInputData input, FeatureHistogram histogram)
}
}

private void SumupManaged(SumupInputData input, FeatureHistogram histogram)
{
unsafe
{
fixed (byte* pData = _data)
{
SumupManagedDense(input, histogram, pData, 8);
}
}
}

public override void Sumup(SumupInputData input, FeatureHistogram histogram) => SumupHandler(input, histogram);
}

Expand All @@ -476,14 +604,14 @@ public Dense4BitIntArray(int len)
: base(len)
{
_data = new byte[(len + 1) / 2]; // Even length = half the bytes. Odd length = half the bytes+0.5.
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public Dense4BitIntArray(int len, IEnumerable<int> values)
: base(len)
{
_data = new byte[(len + 1) / 2];
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);

int currentIndex = 0;
bool upper = true;
Expand All @@ -508,7 +636,7 @@ public Dense4BitIntArray(byte[] buffer, ref int position)
: base(buffer.ToInt(ref position))
{
_data = buffer.ToByteArray(ref position);
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

/// <summary>
Expand Down Expand Up @@ -580,6 +708,17 @@ public void SumupNative(SumupInputData input, FeatureHistogram histogram)
}
}

private void SumupManaged(SumupInputData input, FeatureHistogram histogram)
{
unsafe
{
fixed (byte* pData = _data)
{
SumupManagedDense(input, histogram, pData, 4);
}
}
}

public override void Sumup(SumupInputData input, FeatureHistogram histogram) => SumupHandler(input, histogram);
}

Expand All @@ -596,21 +735,21 @@ public Dense16BitIntArray(int len)
: base(len)
{
_data = new ushort[len];
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public Dense16BitIntArray(int len, IEnumerable<int> values)
: base(len)
{
_data = values.Select(i => (ushort)i).ToArray(len);
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public Dense16BitIntArray(byte[] buffer, ref int position)
: base(buffer.ToInt(ref position))
{
_data = buffer.ToUShortArray(ref position);
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public override unsafe void Callback(Action<IntPtr> callback)
Expand Down Expand Up @@ -668,6 +807,18 @@ public void SumupNative(SumupInputData input, FeatureHistogram histogram)
}
}

private void SumupManaged(SumupInputData input, FeatureHistogram histogram)
{
unsafe
{
fixed (ushort* pData = _data)
{
byte* pDataBytes = (byte*)pData;
SumupManagedDense(input, histogram, pDataBytes, 16);
}
}
}

public override void Sumup(SumupInputData input, FeatureHistogram histogram) => SumupHandler(input, histogram);

}
Expand All @@ -685,21 +836,21 @@ public Dense32BitIntArray(int len)
: base(len)
{
_data = new int[len];
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public Dense32BitIntArray(int len, IEnumerable<int> values)
: base(len)
{
_data = values.ToArray(len);
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public Dense32BitIntArray(byte[] buffer, ref int position)
: base(buffer.ToInt(ref position))
{
_data = buffer.ToIntArray(ref position);
SetupSumupHandler(SumupNative, base.Sumup);
SetupSumupHandler(SumupNative, SumupManaged);
}

public override unsafe void Callback(Action<IntPtr> callback)
Expand Down Expand Up @@ -757,6 +908,18 @@ public void SumupNative(SumupInputData input, FeatureHistogram histogram)
}
}

private void SumupManaged(SumupInputData input, FeatureHistogram histogram)
{
unsafe
{
fixed (int* pData = _data)
{
byte* pDataBytes = (byte*)pData;
SumupManagedDense(input, histogram, pDataBytes, 32);
}
}
}

public override void Sumup(SumupInputData input, FeatureHistogram histogram) => SumupHandler(input, histogram);
}
}
Loading
Loading