Reduce CPU partial factor allocations

This commit is contained in:
Marek Lemańczyk committed 2025-10-10 00:07:39 +02:00
1 parent 74e9e88977
commit 8054f9fe59
1 file changed
+160 -45
+160 -45
View File
@@ -95,6 +95,15 @@ internal static partial class PrimeOrderCalculator
[ThreadStatic]
private static bool s_debugLoggingEnabled;
[ThreadStatic]
private static List<ulong>? s_partialFactorPending;
[ThreadStatic]
private static Stack<ulong>? s_partialFactorCompositeStack;
[ThreadStatic]
private static Dictionary<ulong, int>? s_partialFactorCounts;
private static bool IsGpuHeuristicDevice => s_pow2ModeInitialized && s_deviceMode == PrimeOrderHeuristicDevice.Gpu;
private readonly struct Pow2ModeScope : IDisposable
{
@@ -189,7 +198,7 @@ internal static partial class PrimeOrderCalculator
using var debugScope = UseDebugLogging(!IsGpuHeuristicDevice);
// DebugLog("Partial factoring φ(p)");
PartialFactorResult phiFactors = PartialFactor(phi, config);
using PartialFactorResult phiFactors = PartialFactor(phi, config);
if (phiFactors.Factors is null)
{
@@ -353,6 +362,8 @@ internal static partial class PrimeOrderCalculator
// TODO: Do we do partial factoring of order multiple times?
PartialFactorResult factorization = PartialFactor(order, config);
try
{
if (factorization.Factors is null)
{
return false;
@@ -362,7 +373,6 @@ internal static partial class PrimeOrderCalculator
{
if (factorization.Cofactor <= 1UL)
{
// DebugLog("Cofactor <= 1. No factors found");
return false;
}
@@ -375,7 +385,9 @@ internal static partial class PrimeOrderCalculator
}
// DebugLog("Adding cofactor as prime factor");
factorization = factorization.WithAdditionalPrime(factorization.Cofactor);
PartialFactorResult extended = factorization.WithAdditionalPrime(factorization.Cofactor);
factorization.Dispose();
factorization = extended;
}
ReadOnlySpan<FactorEntry> span = factorization.Factors;
@@ -402,6 +414,11 @@ internal static partial class PrimeOrderCalculator
return true;
}
finally
{
factorization.Dispose();
}
}
private static bool TryHeuristicFinish(
ulong prime,
@@ -421,6 +438,8 @@ internal static partial class PrimeOrderCalculator
// TODO: Do we do partial factoring of order multiple times?
// DebugLog("Trying heuristic. Partial factoring order");
PartialFactorResult orderFactors = PartialFactor(order, config);
try
{
if (orderFactors.Factors is null)
{
return false;
@@ -440,7 +459,9 @@ internal static partial class PrimeOrderCalculator
return false;
}
orderFactors = orderFactors.WithAdditionalPrime(orderFactors.Cofactor);
PartialFactorResult extended = orderFactors.WithAdditionalPrime(orderFactors.Cofactor);
orderFactors.Dispose();
orderFactors = extended;
}
int capacity = config.MaxPowChecks <= 0 ? 64 : config.MaxPowChecks * 4;
@@ -564,6 +585,11 @@ internal static partial class PrimeOrderCalculator
// DebugLog("No candidate confirmed");
return false;
}
finally
{
orderFactors.Dispose();
}
}
private static void SortCandidates(ulong prime, ulong? previousOrder, List<ulong> candidates)
{
@@ -735,6 +761,8 @@ internal static partial class PrimeOrderCalculator
private static bool TryConfirmCandidate(ulong prime, ulong candidate, in MontgomeryDivisorData divisorData, PrimeOrderSearchConfig config, ref int powUsed, int powBudget)
{
PartialFactorResult factorization = PartialFactor(candidate, config);
try
{
if (factorization.Factors is null)
{
return false;
@@ -754,7 +782,9 @@ internal static partial class PrimeOrderCalculator
return false;
}
factorization = factorization.WithAdditionalPrime(factorization.Cofactor);
PartialFactorResult extended = factorization.WithAdditionalPrime(factorization.Cofactor);
factorization.Dispose();
factorization = extended;
}
ReadOnlySpan<FactorEntry> span = factorization.Factors;
@@ -786,6 +816,11 @@ internal static partial class PrimeOrderCalculator
return true;
}
finally
{
factorization.Dispose();
}
}
private static bool Pow2EqualsOne(ulong exponent, ulong prime, in MontgomeryDivisorData divisorData)
{
@@ -905,6 +940,56 @@ internal static partial class PrimeOrderCalculator
}
}
private static List<ulong> AcquirePendingCompositeList(int capacityHint)
{
List<ulong>? list = s_partialFactorPending;
if (list is null)
{
list = new List<ulong>(capacityHint);
s_partialFactorPending = list;
return list;
}
list.Clear();
if (list.Capacity < capacityHint)
{
list.Capacity = capacityHint;
}
return list;
}
private static Stack<ulong> AcquireCompositeStack(int capacityHint)
{
Stack<ulong>? stack = s_partialFactorCompositeStack;
if (stack is null)
{
stack = new Stack<ulong>(capacityHint);
s_partialFactorCompositeStack = stack;
return stack;
}
stack.Clear();
stack.EnsureCapacity(capacityHint);
return stack;
}
private static Dictionary<ulong, int> AcquireFactorCountDictionary(int capacityHint)
{
Dictionary<ulong, int>? dictionary = s_partialFactorCounts;
int adjustedCapacity = Math.Max(capacityHint, 8);
if (dictionary is null)
{
dictionary = new Dictionary<ulong, int>(adjustedCapacity);
s_partialFactorCounts = dictionary;
return dictionary;
}
dictionary.Clear();
dictionary.EnsureCapacity(adjustedCapacity);
return dictionary;
}
private static PartialFactorResult PartialFactor(ulong value, PrimeOrderSearchConfig config)
{
if (value <= 1UL)
@@ -923,18 +1008,19 @@ internal static partial class PrimeOrderCalculator
bool useDictionary = false;
uint limit = config.SmallFactorLimit == 0 ? uint.MaxValue : config.SmallFactorLimit;
ulong remaining;
ulong remaining = value;
bool gpuFactored = false;
if (IsGpuHeuristicDevice)
{
gpuFactored = PrimeOrderGpuHeuristics.TryPartialFactor(value, limit, primeSlots, exponentSlots, out factorCount, out remaining, out _);
}
else
{
remaining = value;
}
List<ulong> pending = AcquirePendingCompositeList(2);
Stack<ulong>? compositeStack = null;
try
{
if (!gpuFactored)
{
primeSlots.Clear();
@@ -942,12 +1028,11 @@ internal static partial class PrimeOrderCalculator
if (!TryPopulateSmallPrimeFactorsCpu(value, limit, primeSlots, exponentSlots, out factorCount, out remaining))
{
useDictionary = true;
counts = new Dictionary<ulong, int>(capacity: 8);
counts = AcquireFactorCountDictionary(Math.Max(factorCount, 8));
remaining = PopulateSmallPrimeFactorsCpu(value, limit, counts);
}
}
List<ulong> pending = new();
if (remaining > 1UL)
{
pending.Add(remaining);
@@ -955,16 +1040,15 @@ internal static partial class PrimeOrderCalculator
if (config.PollardRhoMilliseconds > 0 && pending.Count > 0)
{
// DebugLog("Processing pending composites with Pollard's Rho");
Stopwatch stopwatch = Stopwatch.StartNew();
long budgetTicks = TimeSpan.FromMilliseconds(config.PollardRhoMilliseconds).Ticks;
Stack<ulong> stack = new();
stack.Push(remaining);
compositeStack = AcquireCompositeStack(Math.Max(pending.Count * 2, 4));
compositeStack.Push(remaining);
pending.Clear();
while (stack.Count > 0)
while (compositeStack.Count > 0)
{
ulong composite = stack.Pop();
ulong composite = compositeStack.Pop();
if (composite <= 1UL)
{
continue;
@@ -990,8 +1074,8 @@ internal static partial class PrimeOrderCalculator
}
ulong quotient = composite / factor;
stack.Push(factor);
stack.Push(quotient);
compositeStack.Push(factor);
compositeStack.Push(quotient);
}
}
@@ -1015,7 +1099,6 @@ internal static partial class PrimeOrderCalculator
{
if ((counts is null || counts.Count == 0) && cofactor == value)
{
// DebugLog("cofactor is the same as value, no factors found");
return new PartialFactorResult(null, value, false, 0);
}
@@ -1026,7 +1109,6 @@ internal static partial class PrimeOrderCalculator
FactorEntry[] factors = ArrayPool<FactorEntry>.Shared.Rent(counts.Count);
int index = 0;
// DebugLog(() => $"Collecting {counts.Count} prime factors");
foreach (KeyValuePair<ulong, int> entry in counts)
{
factors[index] = new FactorEntry(entry.Key, entry.Value);
@@ -1036,12 +1118,8 @@ internal static partial class PrimeOrderCalculator
Span<FactorEntry> span = factors.AsSpan(0, index);
span.Sort(static (a, b) => a.Value.CompareTo(b.Value));
FactorEntry[] resultArray = new FactorEntry[index];
span.CopyTo(resultArray);
ArrayPool<FactorEntry>.Shared.Return(factors, clearArray: false);
bool fullyFactored = cofactor == 1UL;
return new PartialFactorResult(resultArray, cofactor, fullyFactored, index);
return new PartialFactorResult(factors, cofactor, fullyFactored, index, pooled: true);
}
int actualCount = 0;
@@ -1057,16 +1135,14 @@ internal static partial class PrimeOrderCalculator
{
if (cofactor == value)
{
// DebugLog("cofactor is the same as value, no factors found");
return new PartialFactorResult(null, value, false, 0);
}
return new PartialFactorResult(null, cofactor, cofactor == 1UL, 0);
}
FactorEntry[] array = new FactorEntry[actualCount];
FactorEntry[] array = ArrayPool<FactorEntry>.Shared.Rent(actualCount);
int arrayIndex = 0;
// DebugLog(() => $"Collecting {actualCount} prime factors");
for (int i = 0; i < factorCount; i++)
{
ulong primeValue = primeSlots[i];
@@ -1080,9 +1156,17 @@ internal static partial class PrimeOrderCalculator
arrayIndex++;
}
Array.Sort(array, static (a, b) => a.Value.CompareTo(b.Value));
Span<FactorEntry> arraySpan = array.AsSpan(0, arrayIndex);
arraySpan.Sort(static (a, b) => a.Value.CompareTo(b.Value));
bool fullyFactoredArray = cofactor == 1UL;
return new PartialFactorResult(array, cofactor, fullyFactoredArray, array.Length);
return new PartialFactorResult(array, cofactor, fullyFactoredArray, arrayIndex, pooled: true);
}
finally
{
pending.Clear();
counts?.Clear();
compositeStack?.Clear();
}
}
private const int GpuSmallPrimeFactorSlots = 64;
@@ -1374,7 +1458,7 @@ internal static partial class PrimeOrderCalculator
return;
}
counts ??= new Dictionary<ulong, int>(Math.Max(count, 8));
counts ??= AcquireFactorCountDictionary(Math.Max(count, 8));
CopyFactorsToDictionary(primes, exponents, count, counts);
count = 0;
useDictionary = true;
@@ -1505,27 +1589,58 @@ internal static partial class PrimeOrderCalculator
public readonly int Exponent = exponent;
}
private readonly struct PartialFactorResult(FactorEntry[]? factors, ulong cofactor, bool fullyFactored, int count)
private struct PartialFactorResult : IDisposable
{
public readonly FactorEntry[]? Factors = factors;
public readonly ulong Cofactor = cofactor;
public readonly bool FullyFactored = fullyFactored;
public readonly int Count = count;
public readonly static PartialFactorResult Empty = new(null, 1UL, true, 0);
private FactorEntry[]? _factors;
private readonly bool _pooled;
public PartialFactorResult(FactorEntry[]? factors, ulong cofactor, bool fullyFactored, int count, bool pooled = false)
{
_factors = factors;
Cofactor = cofactor;
FullyFactored = fullyFactored;
Count = count;
_pooled = pooled && factors is not null;
}
public FactorEntry[]? Factors => _factors;
public ulong Cofactor { get; }
public bool FullyFactored { get; }
public int Count { get; }
public static PartialFactorResult Empty => new(null, 1UL, true, 0);
public void Dispose()
{
FactorEntry[]? factors = _factors;
if (factors is null || !_pooled)
{
return;
}
_factors = null;
ArrayPool<FactorEntry>.Shared.Return(factors, clearArray: false);
}
public PartialFactorResult WithAdditionalPrime(ulong prime)
{
if (Factors is null)
if (_factors is null || Count == 0)
{
FactorEntry[] local = [new FactorEntry(prime, 1)];
return new PartialFactorResult(local, 1UL, true, 1);
FactorEntry[] local = ArrayPool<FactorEntry>.Shared.Rent(1);
local[0] = new FactorEntry(prime, 1);
return new PartialFactorResult(local, 1UL, true, 1, pooled: true);
}
FactorEntry[] extended = new FactorEntry[Count + 1];
Array.Copy(Factors, 0, extended, 0, Count);
FactorEntry[] extended = ArrayPool<FactorEntry>.Shared.Rent(Count + 1);
Array.Copy(_factors, 0, extended, 0, Count);
extended[Count] = new FactorEntry(prime, 1);
Array.Sort(extended, static (a, b) => a.Value.CompareTo(b.Value));
return new PartialFactorResult(extended, 1UL, true, Count + 1);
Span<FactorEntry> span = extended.AsSpan(0, Count + 1);
span.Sort(static (a, b) => a.Value.CompareTo(b.Value));
return new PartialFactorResult(extended, 1UL, true, Count + 1, pooled: true);
}
}
}