Various optimizations

This commit is contained in:
mlemanczyk committed 2025-10-17 05:07:11 +02:00
1 parent 91dd373f43
commit d30a915890
1 file changed
+45 -71
+45 -71
View File
@@ -241,20 +241,20 @@ public static class ULongExtensions
// faster, while GPU code can keep using GpuUInt128.MulHigh.
ulong xLow = (uint)x;
ulong xHigh = x >> 32;
ulong yLow = (uint)y;
ulong result = (uint)y;
ulong yHigh = y >> 32;
ulong w1 = xLow * yHigh;
ulong w2 = xHigh * yLow;
ulong w3 = xLow * yLow;
ulong w2 = xHigh * result;
xLow *= result;
// Matching the layout used in GpuUInt128.MulHigh: introducing the
// intermediate result looks like one extra store, but it lets RyuJIT keep
// the accumulated high word entirely in registers. Without this explicit
// local the JIT spills the partial sum, which is where the performance
// regression in the benchmarks came from.
ulong result = (xHigh * yHigh) + (w1 >> 32) + (w2 >> 32);
result += ((w3 >> 32) + (uint)w1 + (uint)w2) >> 32;
result = (xHigh * yHigh) + (w1 >> 32) + (w2 >> 32);
result += ((xLow >> 32) + (uint)w1 + (uint)w2) >> 32;
return result;
}
@@ -316,9 +316,9 @@ public static class ULongExtensions
{
ulong tLow = unchecked(a * b);
ulong m = unchecked(tLow * nPrime);
ulong mTimesModulusLow = unchecked(m * modulus);
ulong result = unchecked(m * modulus);
ulong result = unchecked(a.MulHigh(b) + m.MulHigh(modulus) + (unchecked(tLow + mTimesModulusLow) < tLow ? 1UL : 0UL));
result = unchecked(a.MulHigh(b) + m.MulHigh(modulus) + (unchecked(tLow + result) < tLow ? 1UL : 0UL));
if (result >= modulus)
{
result -= modulus;
@@ -424,10 +424,11 @@ public static class ULongExtensions
internal static ulong Pow2MontgomeryModWindowedKernel(ulong exponent, in MontgomeryDivisorData divisor, bool keepMontgomery)
{
ulong modulus = divisor.Modulus;
if (exponent == 0UL)
{
return keepMontgomery ? divisor.MontgomeryOne : 1UL % modulus;
}
// This will never happen in production code
// if (exponent == 0UL)
// {
// return keepMontgomery ? divisor.MontgomeryOne : 1UL % modulus;
// }
if (exponent <= Pow2WindowFallbackThreshold)
{
@@ -450,26 +451,27 @@ public static class ULongExtensions
}
int windowStart = index - windowSize + 1;
if (windowStart < 0)
{
windowStart = 0;
}
// This can never happen in production code
// if (windowStart < 0)
// {
// windowStart = 0;
// }
while (((exponent >> windowStart) & 1UL) == 0UL)
{
windowStart++;
}
int windowLength = index - windowStart + 1;
for (int square = 0; square < windowLength; square++)
bitLength = index - windowStart + 1;
for (int square = 0; square < bitLength; square++)
{
result = result.MontgomeryMultiply(result, modulus, nPrime);
}
ulong mask = (1UL << windowLength) - 1UL;
ulong windowValue = (exponent >> windowStart) & mask;
ulong multiplier = ComputeMontgomeryOddPower(windowValue, divisor, modulus, nPrime);
result = result.MontgomeryMultiply(multiplier, modulus, nPrime);
ulong windowValue = (1UL << bitLength) - 1UL;
windowValue = (exponent >> windowStart) & windowValue;
windowValue = ComputeMontgomeryOddPower(windowValue, divisor, modulus, nPrime);
result = result.MontgomeryMultiply(windowValue, modulus, nPrime);
index = windowStart - 1;
}
@@ -527,8 +529,6 @@ public static class ULongExtensions
}
var execution = lease.EnterExecutionScope();
try
{
Accelerator accelerator = lease.Accelerator;
if (accelerator.AcceleratorType == AcceleratorType.CPU)
{
@@ -542,11 +542,12 @@ public static class ULongExtensions
return launcher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, byte, ArrayView1D<ulong, Stride1D.Dense>>>();
});
using var exponentBuffer = accelerator.Allocate1D<ulong>(1);
using var resultBuffer = accelerator.Allocate1D<ulong>(1);
var exponentBuffer = accelerator.Allocate1D<ulong>(1);
var resultBuffer = accelerator.Allocate1D<ulong>(1);
exponentBuffer.View.CopyFromCPU(ref exponent, 1);
resultBuffer.MemSetToZero();
// We could clear the buffer, but it'll be overwritten anyway
// resultBuffer.MemSetToZero();
byte keepFlag = keepMontgomery ? (byte)1 : (byte)0;
AcceleratorStream stream = lease.Stream;
@@ -554,21 +555,12 @@ public static class ULongExtensions
stream.Synchronize();
resultBuffer.View.CopyToCPU(ref result, 1);
return true;
}
catch (AcceleratorException)
{
return false;
}
catch (NotSupportedException)
{
return false;
}
finally
{
resultBuffer.Dispose();
exponentBuffer.Dispose();
execution.Dispose();
lease.Dispose();
}
return true;
}
private static void Pow2MontgomeryKernel(Index1D index, ArrayView1D<ulong, Stride1D.Dense> exponents, MontgomeryDivisorData divisor, byte keepMontgomery, ArrayView1D<ulong, Stride1D.Dense> results)
@@ -608,43 +600,25 @@ public static class ULongExtensions
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static int GetPortableBitLength(ulong value)
{
if (value == 0UL)
{
return 0;
}
// This will never happen in production code
// if (value == 0UL)
// {
// return 0;
// }
return 64 - XMath.LeadingZeroCount(value);
}
private static int GetWindowSize(int bitLength)
private static int GetWindowSize(int bitLength) => bitLength switch
{
if (bitLength <= 6)
{
return Math.Max(bitLength, 1);
}
if (bitLength <= 23)
{
return 4;
}
if (bitLength <= 79)
{
return 5;
}
if (bitLength <= 239)
{
return 6;
}
if (bitLength <= 671)
{
return 7;
}
return Pow2WindowSize;
}
<= 6 => bitLength,// bitLength is always >= 1 in production code
// return Math.Max(bitLength, 1);
<= 23 => 4,
<= 79 => 5,
<= 239 => 6,
<= 671 => 7,
_ => Pow2WindowSize,
};
[MethodImpl(MethodImplOptions.AggressiveInlining)]
private static void InitializeMontgomeryOddPowersCpu(in MontgomeryDivisorData divisor, ulong modulus, ulong nPrime, Span<ulong> oddPowers)