Refactor heuristic pooling to cut GC pressure

PerfectNumbers.Core/PrimeOrderCalculator.cs

@@ -48,17 +48,9 @@ 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;
 
-	[ThreadStatic]
-	private static List<ulong>? s_heuristicCandidateList;
 
 	private static bool IsGpuHeuristicDevice => s_pow2ModeInitialized && s_deviceMode == PrimeOrderHeuristicDevice.Gpu;
 
@@ -406,13 +398,13 @@ internal static partial class PrimeOrderCalculator
 			}
 
 			int capacity = config.MaxPowChecks <= 0 ? 64 : config.MaxPowChecks * 4;
-			List<ulong> candidates = AcquireHeuristicCandidateList(capacity);
+			List<ulong> candidates = ThreadStaticPools.RentUlongList(capacity);
 			FactorEntry[] factorArray = orderFactors.Factors!;
 			// DebugLog("Building candidates list");
 			BuildCandidates(order, factorArray, orderFactors.Count, candidates, capacity);
 			if (candidates.Count == 0)
 			{
-				candidates.Clear();
+				ThreadStaticPools.ReturnUlongList(candidates);
 				return false;
 			}
 
@@ -513,7 +505,7 @@ internal static partial class PrimeOrderCalculator
 						pool.Return(gpuPool, clearArray: false);
 					}
 
-					candidates.Clear();
+					ThreadStaticPools.ReturnUlongList(candidates);
 					result = candidate;
 					return true;
 				}
@@ -526,7 +518,7 @@ internal static partial class PrimeOrderCalculator
 				index += batchSize;
 			}
 
-			candidates.Clear();
+			ThreadStaticPools.ReturnUlongList(candidates);
                         // DebugLog("No candidate confirmed");
 			return false;
 		}
@@ -781,70 +773,6 @@ internal static partial class PrimeOrderCalculator
 		return exponent.Pow2MontgomeryModWindowedCpu(divisorData, keepMontgomery: false) == 1UL;
 	}
 
-	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;
-		if (dictionary is null)
-		{
-			dictionary = new Dictionary<ulong, int>(capacityHint);
-			s_partialFactorCounts = dictionary;
-			return dictionary;
-		}
-
-		dictionary.Clear();
-		dictionary.EnsureCapacity(capacityHint);
-		return dictionary;
-	}
-
-	private static List<ulong> AcquireHeuristicCandidateList(int capacityHint)
-	{
-		List<ulong>? list = s_heuristicCandidateList;
-		if (list is null)
-		{
-			list = new List<ulong>(capacityHint);
-			s_heuristicCandidateList = list;
-			return list;
-		}
-
-		list.Clear();
-		list.EnsureCapacity(capacityHint);
-		return list;
-	}
-
 	private static PartialFactorResult PartialFactor(ulong value, in PrimeOrderSearchConfig config)
         {
                 if (value <= 1UL)
@@ -871,11 +799,10 @@ internal static partial class PrimeOrderCalculator
                         gpuFactored = PrimeOrderGpuHeuristics.TryPartialFactor(value, limit, primeSlots, exponentSlots, out factorCount, out remaining, out _);
                 }
 
-		List<ulong> pending = AcquirePendingCompositeList(2);
+                List<ulong> pending = ThreadStaticPools.RentUlongList(2);
                 Stack<ulong>? compositeStack = null;
+                PartialFactorResult result;
 
-		try
-		{
                 if (!gpuFactored)
                 {
                         primeSlots.Clear();
@@ -883,7 +810,7 @@ internal static partial class PrimeOrderCalculator
                         if (!TryPopulateSmallPrimeFactorsCpu(value, limit, primeSlots, exponentSlots, out factorCount, out remaining))
                         {
                                 useDictionary = true;
-					counts = AcquireFactorCountDictionary(Math.Max(factorCount, 8));
+                                counts = ThreadStaticPools.RentUlongIntDictionary(Math.Max(factorCount, 8));
                                 remaining = PopulateSmallPrimeFactorsCpu(value, limit, counts);
                         }
                 }
@@ -896,11 +823,11 @@ internal static partial class PrimeOrderCalculator
                 if (config.PollardRhoMilliseconds > 0 && pending.Count > 0)
                 {
                         long deadlineTimestamp = CreateDeadlineTimestamp(config.PollardRhoMilliseconds);
-				compositeStack = AcquireCompositeStack(Math.Max(pending.Count * 2, 4));
+                        compositeStack = ThreadStaticPools.RentUlongStack(Math.Max(pending.Count * 2, 4));
                         compositeStack.Push(remaining);
                         pending.Clear();
 
-				long timestamp = 0L; // reused for deadline checks.
+                        long timestamp = 0L;
                         while (compositeStack.Count > 0)
                         {
                                 ulong composite = compositeStack.Pop();
@@ -956,12 +883,14 @@ internal static partial class PrimeOrderCalculator
                 {
                         if ((counts is null || counts.Count == 0) && cofactor == value)
                         {
-					return PartialFactorResult.Rent(null, value, false, 0);
+                                result = PartialFactorResult.Rent(null, value, false, 0);
+                                goto ReturnResult;
                         }
 
                         if (counts is null || counts.Count == 0)
                         {
-					return PartialFactorResult.Rent(null, cofactor, cofactor == 1UL, 0);
+                                result = PartialFactorResult.Rent(null, cofactor, cofactor == 1UL, 0);
+                                goto ReturnResult;
                         }
 
                         FactorEntry[] factors = pool.Rent(counts.Count);
@@ -976,7 +905,8 @@ internal static partial class PrimeOrderCalculator
                         span.Sort(static (a, b) => a.Value.CompareTo(b.Value));
 
                         bool fullyFactored = cofactor == 1UL;
-				return PartialFactorResult.Rent(factors, cofactor, fullyFactored, index);
+                        result = PartialFactorResult.Rent(factors, cofactor, fullyFactored, index);
+                        goto ReturnResult;
                 }
 
                 int actualCount = 0;
@@ -992,10 +922,12 @@ internal static partial class PrimeOrderCalculator
                 {
                         if (cofactor == value)
                         {
-					return PartialFactorResult.Rent(null, value, false, 0);
+                                result = PartialFactorResult.Rent(null, value, false, 0);
+                                goto ReturnResult;
                         }
 
-				return PartialFactorResult.Rent(null, cofactor, cofactor == 1UL, 0);
+                        result = PartialFactorResult.Rent(null, cofactor, cofactor == 1UL, 0);
+                        goto ReturnResult;
                 }
 
                 FactorEntry[] array = pool.Rent(actualCount);
@@ -1016,14 +948,21 @@ internal static partial class PrimeOrderCalculator
                 Span<FactorEntry> arraySpan = array.AsSpan(0, arrayIndex);
                 arraySpan.Sort(static (a, b) => a.Value.CompareTo(b.Value));
                 bool fullyFactoredArray = cofactor == 1UL;
-			return PartialFactorResult.Rent(array, cofactor, fullyFactoredArray, arrayIndex);
+                result = PartialFactorResult.Rent(array, cofactor, fullyFactoredArray, arrayIndex);
-		}
+
-		finally
+        ReturnResult:
+                if (counts is not null)
                 {
-			pending.Clear();
+                        ThreadStaticPools.ReturnUlongIntDictionary(counts);
-			counts?.Clear();
-			compositeStack?.Clear();
                 }
+
+                ThreadStaticPools.ReturnUlongList(pending);
+                if (compositeStack is not null)
+                {
+                        ThreadStaticPools.ReturnUlongStack(compositeStack);
+                }
+
+                return result;
         }
 
         private const int GpuSmallPrimeFactorSlots = 64;
@@ -1340,7 +1279,7 @@ internal static partial class PrimeOrderCalculator
 			return;
 		}
 
-		counts ??= AcquireFactorCountDictionary(Math.Max(count, 8));
+		counts ??= ThreadStaticPools.RentUlongIntDictionary(Math.Max(count, 8));
 		CopyFactorsToDictionary(primes, exponents, count, counts);
 		count = 0;
 		useDictionary = true;

PerfectNumbers.Core/ThreadStaticPools.cs

@@ -79,6 +79,130 @@ namespace PerfectNumbers.Core
             }
         }
 
+
+        [ThreadStatic]
+        private static List<List<ulong>>? _ulongListPool;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static List<ulong> RentUlongList(int capacityHint)
+        {
+            List<List<ulong>>? pool = _ulongListPool;
+            if (pool is not null && pool.Count > 0)
+            {
+                int lastIndex = pool.Count - 1;
+                List<ulong> list = pool[lastIndex];
+                pool.RemoveAt(lastIndex);
+                if (list.Count > 0)
+                {
+                    list.Clear();
+                }
+
+                if (list.Capacity < capacityHint)
+                {
+                    list.Capacity = capacityHint;
+                }
+
+                return list;
+            }
+
+            return new List<ulong>(capacityHint);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void ReturnUlongList(List<ulong> list)
+        {
+            List<List<ulong>>? pool = _ulongListPool;
+            if (pool is null)
+            {
+                pool = new List<List<ulong>>(4);
+                _ulongListPool = pool;
+            }
+
+            pool.Add(list);
+        }
+
+        [ThreadStatic]
+        private static List<Stack<ulong>>? _ulongStackPool;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static Stack<ulong> RentUlongStack(int capacityHint)
+        {
+            List<Stack<ulong>>? pool = _ulongStackPool;
+            if (pool is not null && pool.Count > 0)
+            {
+                int lastIndex = pool.Count - 1;
+                Stack<ulong> stack = pool[lastIndex];
+                pool.RemoveAt(lastIndex);
+                if (stack.Count > 0)
+                {
+                    stack.Clear();
+                }
+
+                if (stack.Count < capacityHint)
+                {
+                    stack.EnsureCapacity(capacityHint);
+                }
+
+                return stack;
+            }
+
+            return new Stack<ulong>(capacityHint);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void ReturnUlongStack(Stack<ulong> stack)
+        {
+            List<Stack<ulong>>? pool = _ulongStackPool;
+            if (pool is null)
+            {
+                pool = new List<Stack<ulong>>(4);
+                _ulongStackPool = pool;
+            }
+
+            pool.Add(stack);
+        }
+
+        [ThreadStatic]
+        private static List<Dictionary<ulong, int>>? _ulongIntDictionaryPool;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static Dictionary<ulong, int> RentUlongIntDictionary(int capacityHint)
+        {
+            List<Dictionary<ulong, int>>? pool = _ulongIntDictionaryPool;
+            if (pool is not null && pool.Count > 0)
+            {
+                int lastIndex = pool.Count - 1;
+                Dictionary<ulong, int> dictionary = pool[lastIndex];
+                pool.RemoveAt(lastIndex);
+                if (dictionary.Count > 0)
+                {
+                    dictionary.Clear();
+                }
+
+                if (dictionary.Count < capacityHint)
+                {
+                    dictionary.EnsureCapacity(capacityHint);
+                }
+
+                return dictionary;
+            }
+
+            return new Dictionary<ulong, int>(capacityHint);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void ReturnUlongIntDictionary(Dictionary<ulong, int> dictionary)
+        {
+            List<Dictionary<ulong, int>>? pool = _ulongIntDictionaryPool;
+            if (pool is null)
+            {
+                pool = new List<Dictionary<ulong, int>>(4);
+                _ulongIntDictionaryPool = pool;
+            }
+
+            pool.Add(dictionary);
+        }
+
         [ThreadStatic]
         private static Dictionary<ulong, MersenneDivisorCycles.FactorCacheEntry>? _mersenneFactorCacheDictionary;
 

PerfectNumbers.Core/ULongExtensions.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Buffers;
 using System.Numerics;
 using System.Runtime.CompilerServices;
 using ILGPU.Algorithms;
@@ -347,24 +346,7 @@ public static class ULongExtensions
         ulong result = divisor.MontgomeryOne;
         ulong nPrime = divisor.NPrime;
 
-        const int StackOddPowerThreshold = PerfectNumberConstants.PooledArrayThreshold;
+        Span<ulong> oddPowers = stackalloc ulong[oddPowerCount];
-        Span<ulong> oddPowersSpan = oddPowerCount <= StackOddPowerThreshold
-            ? stackalloc ulong[StackOddPowerThreshold]
-            : default;
-
-        Span<ulong> oddPowers;
-        ulong[]? oddPowersArray = null;
-
-        if (!oddPowersSpan.IsEmpty)
-        {
-            oddPowers = oddPowersSpan.Slice(0, oddPowerCount);
-        }
-        else
-        {
-            oddPowersArray = ThreadStaticPools.UlongPool.Rent(oddPowerCount);
-            oddPowers = oddPowersArray.AsSpan(0, oddPowerCount);
-        }
-
         InitializeMontgomeryOddPowersCpu(divisor, modulus, nPrime, oddPowers);
 
         int index = bitLength - 1;
@@ -403,10 +385,6 @@ public static class ULongExtensions
             index = windowStart - 1;
         }
 
-        if (oddPowersArray is not null)
-        {
-            ThreadStaticPools.UlongPool.Return(oddPowersArray, clearArray: false);
-        }
 
         if (keepMontgomery)
         {