Various optimizations to separate non-static data between threads and resolve read/write access issue.

PerfectNumbers.Core/Gpu/GpuContextPool.cs

@@ -4,6 +4,21 @@ using ILGPU.Runtime;
 
 namespace PerfectNumbers.Core.Gpu;
 
+public sealed class GpuContextLease
+{
+	private GpuContextPool.PooledContext? _ctx;
+
+	internal GpuContextLease(GpuContextPool.PooledContext ctx)
+	{
+		_ctx = ctx;
+	}
+
+	public Context Context => _ctx!.Context;
+	public Accelerator Accelerator => _ctx!.Accelerator;
+
+	public void Dispose() => GpuContextPool.Return(_ctx!);
+}
+
 public static class GpuContextPool
 {
 	// The pool intentionally skips pre-loading ProcessEightBitWindows kernels so accelerator initialization
@@ -14,12 +29,10 @@ public static class GpuContextPool
 		public readonly Context Context;
 		public readonly Accelerator Accelerator;
 
-		public PooledContext(KernelContainer kernels)
+		public PooledContext()
 		{
 			Context = SharedGpuContext.Context;
 			Accelerator = SharedGpuContext.Accelerator;
-			GpuStaticTableInitializer.EnsureStaticTables(kernels, Accelerator);
-			kernels.Dispose();
 			// NOTE: Avoid loading/compiling any kernel here to prevent implicit
 			// CL stream/queue creation during accelerator construction.
 			// Some OpenCL drivers are fragile when a queue is created immediately
@@ -29,11 +42,9 @@ public static class GpuContextPool
 		public void Dispose()
 		{
 			// Ensure all cached GPU buffers for this accelerator are released.
-			NttGpuMath.ClearCaches(Accelerator);
 			// Release PrimeTester GPU state for this accelerator (kernel/device primes)
 			// Clear any per-accelerator cached resources to avoid releasing
 			// after the accelerator is destroyed.
-			PrimeTester.ClearGpuCaches(Accelerator);
 
 			// These resources are shared between GPU leases
 			//  Accelerator.Dispose();
@@ -53,7 +64,7 @@ public static class GpuContextPool
 		}
 
 		// Create a new accelerator when the pool does not have one available.
-		return new GpuContextLease(new PooledContext(GpuKernelPool.GetKernels()));
+		return new GpuContextLease(new PooledContext());
 	}
 
 
@@ -71,8 +82,7 @@ public static class GpuContextPool
 		PooledContext context;
 		for (int i = 0; i < toCreate; i++)
 		{
-			KernelContainer kernels = GpuKernelPool.GetKernels();
-			context = new PooledContext(kernels);
+			context = new PooledContext();
 			contexts[i] = context;
 			pool.Enqueue(context);
 		}
@@ -87,41 +97,12 @@ public static class GpuContextPool
 			gpu.Dispose();
 		}
 
-		PrimeTester.DisposeGpuContexts();
 		WarmedGpuContextCount = 0;
 	}
 
-	private static void Return(PooledContext ctx)
+	internal static void Return(PooledContext ctx)
 	{
-		ctx.Accelerator.Synchronize();
+		// ctx.Accelerator.Synchronize();
 		GpuPool.Enqueue(ctx);
 	}
-
-	public struct GpuContextLease
-	{
-		private PooledContext? _ctx;
-		private bool _disposed;
-
-		internal GpuContextLease(PooledContext ctx)
-		{
-			_ctx = ctx;
-			_disposed = false;
-		}
-
-		public Context Context => _ctx!.Context;
-		public Accelerator Accelerator => _ctx!.Accelerator;
-
-		public void Dispose()
-		{
-			if (_disposed)
-			{
-				return;
-			}
-
-			_disposed = true;
-			Return(_ctx!);
-			_ctx = null;
-		}
-	}
-}
-
+}

PerfectNumbers.Core/Gpu/GpuKernelLease.cs

@@ -1,9 +1,6 @@
-using System;
 using System.Collections.Concurrent;
 using ILGPU;
 using ILGPU.Runtime;
-using PerfectNumbers.Core;
-using static PerfectNumbers.Core.Gpu.GpuContextPool;
 
 namespace PerfectNumbers.Core.Gpu;
 
@@ -11,11 +8,13 @@ public sealed class GpuKernelLease
 {
     private static readonly ConcurrentQueue<GpuKernelLease> Pool = new();
 
-    private GpuContextLease? _gpu;
+    private GpuContextLease _gpu;
     public KernelContainer Kernels;
-    private AcceleratorStream? _stream;
+    public AcceleratorStream Stream;
 
-    private GpuKernelLease(KernelContainer kernels)
+#pragma warning disable CS8618 // _gpu and Stream will be set by Rent method
+	private GpuKernelLease(KernelContainer kernels)
+#pragma warning restore CS8618
 	{
 		Kernels = kernels;
 	}
@@ -28,35 +27,19 @@ public sealed class GpuKernelLease
         }
 
         lease._gpu = gpu;
-        lease.Kernels = kernels;
-        lease._stream = null;
+		lease.Kernels = kernels;
+		lease.Stream = gpu.Accelerator.CreateStream();
         return lease;
     }
 
     public Accelerator Accelerator => _gpu?.Accelerator ?? throw new NullReferenceException("GPU context is already released");
 
-    public AcceleratorStream Stream
-    {
-        get
-        {
-            if (_stream is { } stream)
-            {
-                return stream;
-            }
-
-            // It feels that we can combine these into 1 return statement, but that increases the max stack frame size to 3 from 2
-            stream = Accelerator.CreateStream();
-            _stream = stream;
-            return stream;
-        }
-    }
-
     public Action<AcceleratorStream, Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>> OrderKernel
     {
         get
         {
             var accel = Accelerator; // avoid capturing 'this' in lambda
-            return KernelContainer.InitOnce(ref Kernels!.Order, () =>
+            return GpuKernelPool.InitOnce(ref Kernels!.Order, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>(OrderKernels.OrderKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -70,7 +53,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref Kernels!.Pow2Mod, () =>
+            return GpuKernelPool.InitOnce(ref Kernels!.Pow2Mod, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernels.Pow2ModKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -84,7 +67,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref Kernels!.Incremental, () =>
+            return GpuKernelPool.InitOnce(ref Kernels!.Incremental, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -98,7 +81,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref Kernels!.IncrementalOrder, () =>
+            return GpuKernelPool.InitOnce(ref Kernels!.IncrementalOrder, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalOrderKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -112,7 +95,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref Kernels!.Pow2ModOrder, () =>
+            return GpuKernelPool.InitOnce(ref Kernels!.Pow2ModOrder, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernels.Pow2ModOrderKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -121,16 +104,16 @@ public sealed class GpuKernelLease
         }
     }
 
-    public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>> SpecialMaxKernel
+    public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> SpecialMaxKernel
     {
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref Kernels!.SpecialMax, () =>
+            return GpuKernelPool.InitOnce(ref Kernels!.SpecialMax, () =>
             {
-                var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel);
+                var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel);
                 var kernel = KernelUtil.GetKernel(loaded);
-                return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>>>();
+                return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>>();
             });
         }
     }
@@ -140,7 +123,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref Kernels!.SmallPrimeFactor, () =>
+            return GpuKernelPool.InitOnce(ref Kernels!.SmallPrimeFactor, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, uint, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(SmallPrimeFactorKernels.SmallPrimeFactorKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -151,15 +134,10 @@ public sealed class GpuKernelLease
 
     public void Dispose()
     {
-        _stream?.Dispose();
-        _stream = null;
-
-        _gpu?.Dispose();
-        _gpu = null;
-
-		Kernels?.Dispose();
-		GpuPrimeWorkLimiter.Release();
+        Stream.Dispose();
+		_gpu.Dispose();
         Pool.Enqueue(this);
+		GpuPrimeWorkLimiter.Release();
     }
 
 }

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -1,4 +1,3 @@
-using PerfectNumbers.Core;
 using System.Collections.Concurrent;
 using ILGPU;
 using ILGPU.Runtime;
@@ -23,7 +22,8 @@ public readonly struct ResidueAutomatonArgs
 	}
 }
 
-public sealed class KernelContainer
+
+public sealed class KernelContainer : IDisposable
 {
 	// Serializes first-time initialization of kernels/buffers per accelerator.
 	public Action<AcceleratorStream, Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>? Order;
@@ -38,42 +38,37 @@ public sealed class KernelContainer
 	public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
 		ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>? Pow2ModOrder;
 	public Action<AcceleratorStream, Index1D, ulong, uint, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>? SmallPrimeFactor;
-	public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>>? SpecialMax;
+	public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>? SpecialMax;
 
 	// Optional device buffer with small divisor cycles (<= 4M). Index = divisor, value = cycle length.
 	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallCycles;
-	public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimeFactorsPrimes;
-	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimeFactorsSquares;
-	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimeFactorPrimeSlots;
-	public MemoryBuffer1D<int, Stride1D.Dense>? SmallPrimeFactorExponentSlots;
-	public MemoryBuffer1D<int, Stride1D.Dense>? SmallPrimeFactorCountSlot;
-	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimeFactorRemainingSlot;
-	public MemoryBuffer1D<ulong, Stride1D.Dense>? SpecialMaxFactors;
-	public MemoryBuffer1D<ulong, Stride1D.Dense>? SpecialMaxCandidates;
-	public MemoryBuffer1D<ushort, Stride1D.Dense>? SpecialMaxResult;
+
 	public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastOne;
-	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimesPow2LastOne;
 	public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastSeven;
+	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimesPow2LastOne;
 	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimesPow2LastSeven;
 
-	public static T InitOnce<T>(ref T? slot, Func<T> factory) where T : class
+	public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimeFactorsPrimes;
+	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimeFactorsSquares;
+
+	public void Dispose()
 	{
-		var current = Volatile.Read(ref slot);
-		if (current is not null)
-		{
-			return current;
-		}
+		// Order = null;
+		// Incremental = null;
+		// Pow2Mod = null;
+		// IncrementalOrder = null;
+		// Pow2ModOrder = null;
+		// SmallPrimeFactor = null;
+		// SpecialMax = null;
 
-		current = Volatile.Read(ref slot);
-		if (current is null)
-		{
-			current = factory();
-			Volatile.Write(ref slot, current);
-		}
-		return current;
+		// SmallCycles?.Dispose();
+		// SmallPrimesLastOne?.Dispose();
+		// SmallPrimesLastSeven?.Dispose();
+		// SmallPrimesPow2LastOne?.Dispose();
+		// SmallPrimesPow2LastSeven?.Dispose();
+		// SmallPrimeFactorsPrimes?.Dispose();
+		// SmallPrimeFactorsSquares?.Dispose();
 	}
-
-	public void Dispose() => GpuKernelPool.Return(this);
 }
 
 public readonly struct ResiduePrimeViews(
@@ -101,49 +96,30 @@ public readonly struct SmallPrimeFactorTables(
 	public ArrayView1D<ulong, Stride1D.Dense> SquaresView => Squares.View;
 }
 
-public readonly struct SmallPrimeFactorScratch(
-		MemoryBuffer1D<ulong, Stride1D.Dense> primeSlots,
-		MemoryBuffer1D<int, Stride1D.Dense> exponentSlots,
-		MemoryBuffer1D<int, Stride1D.Dense> countSlot,
-		MemoryBuffer1D<ulong, Stride1D.Dense> remainingSlot)
-{
-	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> PrimeSlots = primeSlots;
-	public readonly MemoryBuffer1D<int, Stride1D.Dense> ExponentSlots = exponentSlots;
-	public readonly MemoryBuffer1D<int, Stride1D.Dense> CountSlot = countSlot;
-	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> RemainingSlot = remainingSlot;
-
-	public void Clear()
-	{
-		PrimeSlots.MemSetToZero();
-		ExponentSlots.MemSetToZero();
-		CountSlot.MemSetToZero();
-		RemainingSlot.MemSetToZero();
-	}
-}
-
-public readonly struct SpecialMaxScratch(
-	MemoryBuffer1D<ulong, Stride1D.Dense> factorValues,
-	MemoryBuffer1D<ulong, Stride1D.Dense> candidateValues,
-	MemoryBuffer1D<ushort, Stride1D.Dense> resultSlot)
-{
-	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> FactorValues = factorValues;
-	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> CandidateValues = candidateValues;
-	public readonly MemoryBuffer1D<ushort, Stride1D.Dense> ResultSlot = resultSlot;
-
-	public ArrayView1D<ulong, Stride1D.Dense> FactorsView => FactorValues.View;
-	public ArrayView1D<ulong, Stride1D.Dense> CandidatesView => CandidateValues.View;
-	public ArrayView1D<ushort, Stride1D.Dense> ResultView => ResultSlot.View;
-}
-
 public class GpuKernelPool
 {
-	private static readonly ConcurrentQueue<KernelContainer> KernelCache = new(); // TODO: Replace this concurrent map with a simple accelerator-indexed lookup once kernel launchers are prewarmed during startup so we can drop the thread-safe wrapper entirely.
+	private static readonly ConcurrentDictionary<Accelerator, KernelContainer> _kernels = new();
 
-	public static KernelContainer GetKernels() => KernelCache.TryDequeue(out var container)
-		? container
-		: new KernelContainer();
+	public static T InitOnce<T>(ref T? target, Func<T> valueFactory) where T : class
+	{
+		if (target is { } value)
+		{
+			return value;
+		}
 
-	public static void Return(KernelContainer container) => KernelCache.Enqueue(container);
+		var newValue = valueFactory();
+		return Interlocked.CompareExchange(ref target, newValue, null) ?? newValue;
+	}
+
+	internal static KernelContainer GetKernels(Accelerator accelerator)
+	{
+		return _kernels.GetOrAdd(accelerator, static (accelerator) =>
+		{
+			var kernels = new KernelContainer();
+			PreloadStaticTables(kernels, accelerator);
+			return kernels;
+		});
+	}
 
 	internal static void PreloadStaticTables(KernelContainer kernels, Accelerator accelerator)
 	{
@@ -215,10 +191,10 @@ public class GpuKernelPool
 			return new SmallPrimeFactorTables(primeBuffer, squareBuffer, (int)primeBuffer.Length);
 		}
 
-		if (kernels.SmallPrimeFactorsPrimes is { } existingPrimes && kernels.SmallPrimeFactorsSquares is { } existingSquares)
-		{
-			return new SmallPrimeFactorTables(existingPrimes, existingSquares, (int)existingPrimes.Length);
-		}
+		// if (kernels.SmallPrimeFactorsPrimes is { } existingPrimes && kernels.SmallPrimeFactorsSquares is { } existingSquares)
+		// {
+		// 	return new SmallPrimeFactorTables(existingPrimes, existingSquares, (int)existingPrimes.Length);
+		// }
 
 		var hostPrimes = PrimesGenerator.SmallPrimes;
 		var hostSquares = PrimesGenerator.SmallPrimesPow2;
@@ -234,74 +210,15 @@ public class GpuKernelPool
 		return new SmallPrimeFactorTables(devicePrimes, deviceSquares, hostPrimes.Length);
 	}
 
-	public static SmallPrimeFactorScratch EnsureSmallPrimeFactorScratch(KernelContainer kernels, Accelerator accelerator, int slotCount)
-	{
-		MemoryBuffer1D<ulong, Stride1D.Dense>? primeSlots = kernels.SmallPrimeFactorPrimeSlots;
-		if (primeSlots is null || primeSlots.Length < slotCount)
-		{
-			primeSlots?.Dispose();
-			primeSlots = accelerator.Allocate1D<ulong>(slotCount);
-			kernels.SmallPrimeFactorPrimeSlots = primeSlots;
-		}
 
-		MemoryBuffer1D<int, Stride1D.Dense>? exponentSlots = kernels.SmallPrimeFactorExponentSlots;
-		if (exponentSlots is null || exponentSlots.Length < slotCount)
-		{
-			exponentSlots?.Dispose();
-			exponentSlots = accelerator.Allocate1D<int>(slotCount);
-			kernels.SmallPrimeFactorExponentSlots = exponentSlots;
-		}
 
-		MemoryBuffer1D<int, Stride1D.Dense>? countSlot = kernels.SmallPrimeFactorCountSlot;
-		if (countSlot is null)
-		{
-			countSlot = accelerator.Allocate1D<int>(1);
-			kernels.SmallPrimeFactorCountSlot = countSlot;
-		}
 
-		MemoryBuffer1D<ulong, Stride1D.Dense>? remainingSlot = kernels.SmallPrimeFactorRemainingSlot;
-		if (remainingSlot is null)
-		{
-			remainingSlot = accelerator.Allocate1D<ulong>(1);
-			kernels.SmallPrimeFactorRemainingSlot = remainingSlot;
-		}
 
-		return new SmallPrimeFactorScratch(primeSlots!, exponentSlots!, countSlot!, remainingSlot!);
-	}
-
-	public static SpecialMaxScratch EnsureSpecialMaxScratch(KernelContainer kernels, Accelerator accelerator, int factorCapacity)
-	{
-		MemoryBuffer1D<ulong, Stride1D.Dense>? factorValues = kernels.SpecialMaxFactors;
-		if (factorValues is null || factorValues.Length < factorCapacity)
-		{
-			factorValues?.Dispose();
-			factorValues = accelerator.Allocate1D<ulong>(factorCapacity);
-			kernels.SpecialMaxFactors = factorValues;
-		}
-
-		MemoryBuffer1D<ulong, Stride1D.Dense>? candidateValues = kernels.SpecialMaxCandidates;
-		if (candidateValues is null || candidateValues.Length < factorCapacity)
-		{
-			candidateValues?.Dispose();
-			candidateValues = accelerator.Allocate1D<ulong>(factorCapacity);
-			kernels.SpecialMaxCandidates = candidateValues;
-		}
-
-		MemoryBuffer1D<ushort, Stride1D.Dense>? resultSlot = kernels.SpecialMaxResult;
-		if (resultSlot is null)
-		{
-			resultSlot = accelerator.Allocate1D<ushort>(1);
-			kernels.SpecialMaxResult = resultSlot;
-		}
-
-		return new SpecialMaxScratch(factorValues!, candidateValues!, resultSlot!);
-	}
-
-	public static GpuKernelLease GetKernel(bool useGpuOrder)
+	public static GpuKernelLease GetKernel()
 	{
 		GpuPrimeWorkLimiter.Acquire();
 		var gpu = Rent();
-		var kernels = GetKernels();
+		var kernels = GetKernels(gpu.Accelerator);
 		return GpuKernelLease.Rent(gpu, kernels);
 	}
 
@@ -311,12 +228,10 @@ public class GpuKernelPool
 	/// <param name="action">Action to run with (Accelerator, Stream).</param>
 	public static void Run(Action<Accelerator, AcceleratorStream> action)
 	{
-		var lease = GetKernel(useGpuOrder: true);
+		var lease = GetKernel();
 		var accelerator = lease.Accelerator;
-		var stream = accelerator.CreateStream();
+		var stream = lease.Stream;
 		action(accelerator, stream);
-		stream.Dispose();
 		lease.Dispose();
-
 	}
 }

PerfectNumbers.Core/Gpu/GpuScratchBufferPool.cs

@@ -0,0 +1,134 @@
+using System.Collections.Concurrent;
+using ILGPU;
+using ILGPU.Runtime;
+
+namespace PerfectNumbers.Core.Gpu;
+
+public sealed class ScratchBuffer
+{
+	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 ScratchBuffer(Accelerator accelerator, int smallPrimeFactorSlotCount, int specialMaxFactorCapacity)
+	{
+		_accelerator = accelerator;
+
+		// lock (accelerator)
+		// {
+			SmallPrimeFactorPrimeSlots = accelerator.Allocate1D<ulong>(smallPrimeFactorSlotCount);
+			SmallPrimeFactorExponentSlots = accelerator.Allocate1D<int>(smallPrimeFactorSlotCount);
+			SmallPrimeFactorCountSlot = accelerator.Allocate1D<int>(1);
+			SmallPrimeFactorRemainingSlot = accelerator.Allocate1D<ulong>(1);
+
+			SpecialMaxFactors = accelerator.Allocate1D<ulong>(specialMaxFactorCapacity);
+			SpecialMaxCandidates = accelerator.Allocate1D<ulong>(specialMaxFactorCapacity);
+			SpecialMaxResult = accelerator.Allocate1D<ulong>(1);
+		// }
+	}
+
+	public void ResizeSmallPrimeFactorSlots(int newSize)
+	{
+		if (SmallPrimeFactorPrimeSlots.Length < newSize)
+		{
+			SmallPrimeFactorPrimeSlots.Dispose();
+			SmallPrimeFactorExponentSlots.Dispose();
+
+			SmallPrimeFactorPrimeSlots = _accelerator.Allocate1D<ulong>(newSize);
+			SmallPrimeFactorExponentSlots = _accelerator.Allocate1D<int>(newSize);
+		}
+		// else
+		// {
+		// 	SmallPrimeFactorPrimeSlots.MemSetToZero();
+		// 	SmallPrimeFactorExponentSlots.MemSetToZero();
+		// }
+	}
+
+	public void ResizeSpecialMaxFactors(int newSize)
+	{
+		if (SpecialMaxFactors.Length < newSize)
+		{
+			SpecialMaxFactors.Dispose();
+			SpecialMaxCandidates.Dispose();
+
+			SpecialMaxFactors = _accelerator.Allocate1D<ulong>(newSize);
+			SpecialMaxCandidates = _accelerator.Allocate1D<ulong>(newSize);
+		}
+		// else
+		// {
+		// 	SpecialMaxFactors.MemSetToZero();
+		// 	SpecialMaxCandidates.MemSetToZero();
+		// }
+	}
+}
+
+public static class GpuScratchBufferPool
+{
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Queue<ScratchBuffer>>? _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)
+	{
+		const int WarmUpCapacity = 100;
+		// var pool = _pools.GetOrAdd(accelerator, _ => new ConcurrentQueue<ScratchBuffer>());
+		if (_pools is not { })
+		{
+			_pools = new(WarmUpCapacity);
+			WarmUp(WarmUpCapacity, DefaultSmallPrimeFactorSlotCount, DefaultSpecialMaxFactorCapacity);
+		}
+		
+		var pools = _pools;
+		if (!pools.TryGetValue(accelerator, out var pool))
+		{
+			pool = new();
+			pools[accelerator] = pool;
+		}
+		
+		// var pool = _pools.GetOrAdd(accelerator, _ => new ConcurrentQueue<ScratchBuffer>());
+		if (pool.TryDequeue(out var buffer))
+		{
+			if (buffer.SmallPrimeFactorPrimeSlots.Length < smallPrimeFactorSlotCount || buffer.SpecialMaxFactors.Length < specialMaxFactorCapacity)
+			{
+				throw new InvalidOperationException();
+				// lock (accelerator)
+				// {
+					// Console.WriteLine($"Resizing GPU scratch buffer from pool ({buffer.SmallPrimeFactorPrimeSlots.Length} / {smallPrimeFactorSlotCount}), ({buffer.SpecialMaxFactors.Length}/{specialMaxFactorCapacity})");
+					// buffer.ResizeSmallPrimeFactorSlots(smallPrimeFactorSlotCount);
+					// buffer.ResizeSpecialMaxFactors(specialMaxFactorCapacity);
+				// }
+			}
+			return buffer;
+		}
+
+		return new ScratchBuffer(accelerator, smallPrimeFactorSlotCount, specialMaxFactorCapacity);
+	}
+
+	public static void Return(ScratchBuffer buffer)
+	{
+		if (_pools!.TryGetValue(buffer._accelerator, out var pool))
+		{
+			pool.Enqueue(buffer);
+		}
+	}
+
+	public static void WarmUp(int count, int smallPrimeFactorSlots, int specialMaxFactorCapacity)
+	{
+		Accelerator accelerator = SharedGpuContext.Accelerator;
+
+		for (int i = 0; i < count; i++)
+		{
+			var buffer = Rent(accelerator, smallPrimeFactorSlots, specialMaxFactorCapacity);
+			Return(buffer);
+		}
+	}
+}

PerfectNumbers.Core/Gpu/Kernels/PrimeOrderSpecialMaxKernels.cs

@@ -1,6 +1,5 @@
 using ILGPU;
 using ILGPU.Runtime;
-using PerfectNumbers.Core;
 
 namespace PerfectNumbers.Core.Gpu;
 
@@ -13,7 +12,7 @@ internal static partial class PrimeOrderGpuHeuristics
         int factorCount,
         MontgomeryDivisorData divisor,
         ArrayView1D<ulong, Stride1D.Dense> candidates,
-        ArrayView1D<ushort, Stride1D.Dense> resultOut)
+        ArrayView1D<ulong, Stride1D.Dense> resultOut)
     {
         if (index != 0)
         {
@@ -62,7 +61,7 @@ internal static partial class PrimeOrderGpuHeuristics
         var stepper = new ExponentRemainderStepperGpu(divisorPartial);
         if (stepper.InitializeIsUnityGpu(candidates[0]))
         {
-            resultOut[0] = 0;
+			Atomic.Exchange(ref resultOut[0], 0);
             return;
         }
 
@@ -70,11 +69,11 @@ internal static partial class PrimeOrderGpuHeuristics
         {
             if (stepper.ComputeNextIsUnityGpu(candidates[i]))
             {
-                resultOut[0] = 0;
+				Atomic.Exchange(ref resultOut[0], 0);
                 return;
             }
         }
 
-        resultOut[0] = 1;
+        Atomic.Exchange(ref resultOut[0], 1);
     }
 }

PerfectNumbers.Core/Gpu/Kernels/PrimeTesterKernels.cs

@@ -18,19 +18,19 @@ internal static class PrimeTesterKernels
 {
     // GPU kernel: small-prime sieve only. Returns 1 if passes sieve (probable prime), 0 otherwise.
     public static void SmallPrimeSieveKernel(
-        Index1D index,
-        ArrayView<ulong> numbers,
-        ArrayView<uint> smallPrimes,
-        ArrayView<uint> smallPrimesLastOne,
-        ArrayView<uint> smallPrimesLastSeven,
-        ArrayView<uint> smallPrimesLastThree,
-        ArrayView<uint> smallPrimesLastNine,
-        ArrayView<ulong> smallPrimesPow2,
-        ArrayView<ulong> smallPrimesPow2LastOne,
-        ArrayView<ulong> smallPrimesPow2LastSeven,
-        ArrayView<ulong> smallPrimesPow2LastThree,
-        ArrayView<ulong> smallPrimesPow2LastNine,
-        ArrayView<byte> results)
+		Index1D index,
+		ArrayView<ulong> numbers,
+		ArrayView<uint> smallPrimes,
+		ArrayView<uint> smallPrimesLastOne,
+		ArrayView<uint> smallPrimesLastSeven,
+		ArrayView<uint> smallPrimesLastThree,
+		ArrayView<uint> smallPrimesLastNine,
+		ArrayView<ulong> smallPrimesPow2,
+		ArrayView<ulong> smallPrimesPow2LastOne,
+		ArrayView<ulong> smallPrimesPow2LastSeven,
+		ArrayView<ulong> smallPrimesPow2LastThree,
+		ArrayView<ulong> smallPrimesPow2LastNine,
+		ArrayView<byte> results)
     {
         ulong n = numbers[index];
         byte result = 1;

PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs

@@ -1,3 +1,4 @@
+using PerfectNumbers.Core.Gpu;
 using System;
 using System.Buffers;
 using System.Collections.Concurrent;
@@ -25,7 +26,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
 
     private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>>> _kernelCache = new(AcceleratorReferenceComparer.Instance);
     private readonly ConcurrentBag<BatchResources> _resourcePool = new();
-    private readonly ConcurrentBag<GpuContextPool.GpuContextLease> _acceleratorPool = new();
+    private readonly ConcurrentBag<GpuContextLease> _acceleratorPool = new();
     private readonly ConcurrentBag<DivisorScanSession> _sessionPool = new();
 
     private Action<Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>> GetKernel(Accelerator accelerator) =>
@@ -476,7 +477,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
     public sealed class DivisorScanSession : IMersenneNumberDivisorByDivisorTester.IDivisorScanSession
     {
         private readonly MersenneNumberDivisorByDivisorGpuTester _owner;
-        private readonly GpuContextPool.GpuContextLease _lease;
+        private readonly GpuContextLease _lease;
         private readonly Accelerator _accelerator;
         private Action<Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>> _kernel = null!;
         private MemoryBuffer1D<GpuDivisorPartialData, Stride1D.Dense> _divisorBuffer = null!;
@@ -650,7 +651,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         }
     }
 
-    private GpuContextPool.GpuContextLease RentAccelerator()
+    private GpuContextLease RentAccelerator()
     {
         if (_acceleratorPool.TryTake(out var lease))
         {
@@ -660,7 +661,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         return GpuContextPool.Rent();
     }
 
-    private void ReturnAccelerator(GpuContextPool.GpuContextLease lease) => _acceleratorPool.Add(lease);
+    private void ReturnAccelerator(GpuContextLease lease) => _acceleratorPool.Add(lease);
 
     private BatchResources RentBatchResources(Accelerator accelerator, int capacity)
     {
@@ -899,8 +900,4 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
     }
 
     IMersenneNumberDivisorByDivisorTester.IDivisorScanSession IMersenneNumberDivisorByDivisorTester.CreateDivisorSession() => CreateDivisorSession();
-}
-
-
-
-
+}

PerfectNumbers.Core/Gpu/MersenneNumberDivisorGpuTester.cs

@@ -38,19 +38,21 @@ public sealed class MersenneNumberDivisorGpuTester
     public bool IsDivisible(ulong exponent, in ReadOnlyGpuUInt128 divisor)
     {
         var gpu = GpuContextPool.Rent();
-        var accelerator = gpu.Accelerator;
+		var accelerator = gpu.Accelerator;
+		var stream = accelerator.CreateStream();
         var kernel = GetKernel(accelerator);
         var resultBuffer = _resultBuffers.GetOrAdd(accelerator, acc => acc.Allocate1D<byte>(1));
 		// There is no point in clearing this buffer. We always override item [0] and never use it beyond item [0]
         // resultBuffer.MemSetToZero();
         kernel(1, exponent, divisor, resultBuffer.View);
-        accelerator.Synchronize();
+        stream.Synchronize();
         Span<byte> result = stackalloc byte[1];
-        resultBuffer.View.CopyToCPU(ref result[0], 1);
+        resultBuffer.View.CopyToCPU(stream, ref result[0], 1);
         bool divisible = result[0] != 0;
         result[0] = 0;
-        resultBuffer.View.CopyFromCPU(ref result[0], 1);
-        accelerator.Synchronize();
+        resultBuffer.View.CopyFromCPU(stream, ref result[0], 1);
+		stream.Synchronize();
+		stream.Dispose();
         gpu.Dispose();
         return divisible;
     }

PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs

@@ -16,7 +16,7 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
     {
         throw new NotImplementedException($"GPU incremental scanning requires the device cycle heuristics implementation (kernel: {_kernelType}, GPU order: {_useGpuOrder}).");
 
-        var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+        var gpuLease = GpuKernelPool.GetKernel();
         var accelerator = gpuLease.Accelerator;
         var stream = gpuLease.Stream;
         int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy

PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs

@@ -9,7 +9,7 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
 
 	public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime)
 	{
-		var gpuLease = GpuKernelPool.GetKernel(_useGpuOrder);
+		var gpuLease = GpuKernelPool.GetKernel();
 		var accelerator = gpuLease.Accelerator;
 		var stream = gpuLease.Stream;
 		int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy and avoids TDR

PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs

@@ -1,3 +1,4 @@
+using PerfectNumbers.Core.Gpu;
 using System;
 using System.Buffers;
 using System.Collections.Concurrent;
@@ -13,7 +14,7 @@ namespace PerfectNumbers.Core.Gpu;
 public class MersenneNumberResidueGpuTester(bool useGpuOrder)
 {
     private readonly bool _useGpuOrder = useGpuOrder;
-    private readonly ConcurrentBag<GpuContextPool.GpuContextLease> _acceleratorPool = new();
+    private readonly ConcurrentBag<GpuContextLease> _acceleratorPool = new();
     private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>> _pow2ModKernelCache = new(AcceleratorReferenceComparer.Instance);
     private readonly ConcurrentDictionary<Accelerator, ConcurrentBag<ResidueResources>> _resourcePools = new(AcceleratorReferenceComparer.Instance);
 
@@ -35,7 +36,7 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
         var stream = accelerator.CreateStream();
         var kernel = GetPow2ModKernel(accelerator);
 		var resources = RentResources(accelerator, GpuConstants.ScanBatchSize);
-		var gpuKernels = GpuKernelPool.GetKernels();
+		var gpuKernels = GpuKernelPool.GetKernels(accelerator);
 
         var orderBuffer = resources.OrderBuffer;
         ulong[] orderArray = resources.OrderArray;
@@ -69,7 +70,8 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
             kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
                 kernelArgs, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2);
 
-            stream.Synchronize();
+			stream.Synchronize();
+			
             orderBuffer.View.CopyToCPU(ref ordersRef, currentSize);
             if (!Volatile.Read(ref isPrime))
             {
@@ -127,7 +129,7 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
                     loaded.Method);
             });
 
-    private GpuContextPool.GpuContextLease RentAccelerator()
+    private GpuContextLease RentAccelerator()
     {
         if (_acceleratorPool.TryTake(out var lease))
         {
@@ -137,11 +139,11 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
         return GpuContextPool.Rent();
     }
 
-    private void ReturnAccelerator(GpuContextPool.GpuContextLease lease) => _acceleratorPool.Add(lease);
+    private void ReturnAccelerator(GpuContextLease lease) => _acceleratorPool.Add(lease);
 
     private ResidueResources RentResources(Accelerator accelerator, int capacity)
     {
-        var bag = _resourcePools.GetOrAdd(accelerator, static _ => new ConcurrentBag<ResidueResources>());
+        var bag = _resourcePools.GetOrAdd(accelerator, static _ => []);
         while (bag.TryTake(out var resources))
         {
             if (resources.Capacity >= capacity)
@@ -195,4 +197,4 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
     }
 
 
-}
+}

PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs

@@ -198,7 +198,7 @@ internal static partial class PrimeOrderGpuHeuristics
 
         try
         {
-            var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+            var lease = GpuKernelPool.GetKernel();
             Accelerator accelerator = lease.Accelerator;
             AcceleratorStream stream = lease.Stream;
 
@@ -364,7 +364,7 @@ internal static partial class PrimeOrderGpuHeuristics
     {
         order = 0UL;
 
-        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+        var lease = GpuKernelPool.GetKernel();
         Accelerator accelerator = lease.Accelerator;
         AcceleratorStream stream = lease.Stream;
 
@@ -526,7 +526,7 @@ internal static partial class PrimeOrderGpuHeuristics
 
     private static bool TryComputeOnGpu(ReadOnlySpan<ulong> exponents, ulong prime, in MontgomeryDivisorData divisorData, Span<ulong> results)
     {
-        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+        var lease = GpuKernelPool.GetKernel();
         Accelerator accelerator = lease.Accelerator;
         AcceleratorStream stream = lease.Stream;
         var kernel = GetPow2ModKernel(accelerator);
@@ -581,7 +581,7 @@ internal static partial class PrimeOrderGpuHeuristics
 
         GpuUInt128[]? rentedExponents = null;
         GpuUInt128[]? rentedResults = null;
-        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+        var lease = GpuKernelPool.GetKernel();
 
         try
         {

PerfectNumbers.Core/Gpu/SharedGpuContext.cs

@@ -5,7 +5,7 @@ using ILGPU.Runtime;
 
 namespace PerfectNumbers.Core.Gpu;
 
-internal static class SharedGpuContext
+public static class SharedGpuContext
 {
     private static readonly Lazy<Context> s_context = new(() => ILGPU.Context.CreateDefault(), LazyThreadSafetyMode.ExecutionAndPublication);
     private static readonly Lazy<Device> s_device = new(() => s_context.Value.GetPreferredDevice(false), LazyThreadSafetyMode.ExecutionAndPublication);