Performance improvements

PerfectNumbers.Core/Gpu/GpuScratchBufferPool.cs

@@ -1,23 +1,20 @@
-using System.Collections.Concurrent;
 using ILGPU;
 using ILGPU.Runtime;
 
 namespace PerfectNumbers.Core.Gpu;
 
-public sealed class ScratchBuffer
+public sealed class GpuScratchBuffer
 {
-	public MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimeFactorPrimeSlots { get; private set; }
-	public MemoryBuffer1D<int, Stride1D.Dense> SmallPrimeFactorExponentSlots { get; private set; }
-	public MemoryBuffer1D<int, Stride1D.Dense> SmallPrimeFactorCountSlot { get; private set; }
-	public MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimeFactorRemainingSlot { get; private set; }
-
-	public MemoryBuffer1D<ulong, Stride1D.Dense> SpecialMaxFactors { get; private set; }
-	public MemoryBuffer1D<ulong, Stride1D.Dense> SpecialMaxCandidates { get; private set; }
-	public MemoryBuffer1D<ulong, Stride1D.Dense> SpecialMaxResult { get; private set; }
-
 	internal readonly Accelerator _accelerator;
+	public readonly MemoryBuffer1D<int, Stride1D.Dense> SmallPrimeFactorCountSlot;
+	public MemoryBuffer1D<int, Stride1D.Dense> SmallPrimeFactorExponentSlots;
+	public MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimeFactorPrimeSlots;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimeFactorRemainingSlot;
+	public MemoryBuffer1D<ulong, Stride1D.Dense> SpecialMaxCandidates;
+	public MemoryBuffer1D<ulong, Stride1D.Dense> SpecialMaxFactors;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> SpecialMaxResult;
 
-	public ScratchBuffer(Accelerator accelerator, int smallPrimeFactorSlotCount, int specialMaxFactorCapacity)
+	public GpuScratchBuffer(Accelerator accelerator, int smallPrimeFactorSlotCount, int specialMaxFactorCapacity)
 	{
 		_accelerator = accelerator;
 
@@ -63,18 +60,29 @@ public sealed class ScratchBuffer
 			}
 		}
 	}
+
+	internal void Dispose()
+	{
+		SmallPrimeFactorPrimeSlots.Dispose();
+		SmallPrimeFactorExponentSlots.Dispose();
+		SmallPrimeFactorCountSlot.Dispose();
+		SmallPrimeFactorRemainingSlot.Dispose();
+		SpecialMaxFactors.Dispose();
+		SpecialMaxCandidates.Dispose();
+		SpecialMaxResult.Dispose();
+	}
 }
 
 public static class GpuScratchBufferPool
 {
 	[ThreadStatic]
-	private static Dictionary<Accelerator, Queue<ScratchBuffer>>? _pools;
+	private static Dictionary<Accelerator, Queue<GpuScratchBuffer>>? _pools;
 
 	// private static readonly ConcurrentDictionary<Accelerator, ConcurrentQueue<ScratchBuffer>> _pools = new();
 	private const int DefaultSmallPrimeFactorSlotCount = 64; // From PrimeOrderCalculator.Gpu.cs
 	private const int DefaultSpecialMaxFactorCapacity = 1024; // A reasonable default, will be resized if needed
 
-	public static ScratchBuffer Rent(Accelerator accelerator, int smallPrimeFactorSlotCount, int specialMaxFactorCapacity)
+	public static GpuScratchBuffer Rent(Accelerator accelerator, int smallPrimeFactorSlotCount, int specialMaxFactorCapacity)
 	{
 		const int WarmUpCapacity = 100;
 		// var pool = _pools.GetOrAdd(accelerator, _ => new ConcurrentQueue<ScratchBuffer>());
@@ -107,10 +115,10 @@ public static class GpuScratchBufferPool
 			return buffer;
 		}
 
-		return new ScratchBuffer(accelerator, smallPrimeFactorSlotCount, specialMaxFactorCapacity);
+		return new GpuScratchBuffer(accelerator, smallPrimeFactorSlotCount, specialMaxFactorCapacity);
 	}
 
-	public static void Return(ScratchBuffer buffer)
+	public static void Return(GpuScratchBuffer buffer)
 	{
 		if (_pools!.TryGetValue(buffer._accelerator, out var pool))
 		{
@@ -128,4 +136,22 @@ public static class GpuScratchBufferPool
 		// 	Return(buffer);
 		// }
 	}
+
+	public static void DisposeAll()
+	{
+		var pool = _pools;
+		if (pool is not null)
+		{
+			foreach (var queue in pool.Values)
+			{
+				while (queue.TryDequeue(out var buffer))
+				{
+					buffer.Dispose();
+				}
+			}
+
+			pool.Clear();
+			_pools = null;
+		}
+	}
 }