Merge pull request #140 from mlemanczyk/codex/optimize-partialfactor-in-evenperfectbitscanner

Reduce garbage collector pressure in CPU partial factoring

PerfectNumbers.Core/PrimeOrderCalculator.cs

@@ -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();
+                compositeStack = AcquireCompositeStack(Math.Max(pending.Count * 2, 4));
-            stack.Push(remaining);
+                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);
+                    compositeStack.Push(factor);
-                stack.Push(quotient);
+                    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;
+        private FactorEntry[]? _factors;
-        public readonly ulong Cofactor = cofactor;
+        private readonly bool _pooled;
-        public readonly bool FullyFactored = fullyFactored;
+
-        public readonly int Count = count;
+        public PartialFactorResult(FactorEntry[]? factors, ulong cofactor, bool fullyFactored, int count, bool pooled = false)
-        public readonly static PartialFactorResult Empty = new(null, 1UL, true, 0);
+        {
+            _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)];
+                FactorEntry[] local = ArrayPool<FactorEntry>.Shared.Rent(1);
-                return new PartialFactorResult(local, 1UL, true, 1);
+                local[0] = new FactorEntry(prime, 1);
+                return new PartialFactorResult(local, 1UL, true, 1, pooled: true);
             }
 
-            FactorEntry[] extended = new FactorEntry[Count + 1];
+            FactorEntry[] extended = ArrayPool<FactorEntry>.Shared.Rent(Count + 1);
-            Array.Copy(Factors, 0, extended, 0, Count);
+            Array.Copy(_factors, 0, extended, 0, Count);
             extended[Count] = new FactorEntry(prime, 1);
-            Array.Sort(extended, static (a, b) => a.Value.CompareTo(b.Value));
+            Span<FactorEntry> span = extended.AsSpan(0, Count + 1);
-            return new PartialFactorResult(extended, 1UL, true, Count + 1);
+            span.Sort(static (a, b) => a.Value.CompareTo(b.Value));
+            return new PartialFactorResult(extended, 1UL, true, Count + 1, pooled: true);
         }
     }
 }
+