GPU stabilization improvements to prevent calculation crashes

PerfectNumbers.Core/DivisorCycleCache.cs

@@ -15,8 +15,8 @@ public sealed class DivisorCycleCache
 	private const byte ByteZero = 0;
 	private const byte ByteOne = 1;
 	private const int StackBufferThreshold = 128;
-
-	private readonly Action<AcceleratorStream, Index1D, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> _gpuKernel = LoadKernel(SharedGpuContext.Accelerator);
+	private readonly Accelerator _accelerator;
+	private readonly Action<AcceleratorStream, Index1D, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> _gpuKernel;
 	private volatile ulong[] _snapshot;
 	private volatile bool _useGpuGeneration = true;
 	private readonly int _divisorCyclesBatchSize;
@@ -40,6 +40,8 @@ public sealed class DivisorCycleCache
 	private DivisorCycleCache(int divisorCyclesBatchSize)
 	{
 		_divisorCyclesBatchSize = Math.Max(1, divisorCyclesBatchSize);
+		_accelerator = SharedGpuContext.Device.CreateAccelerator(SharedGpuContext.Context);
+		_gpuKernel = LoadKernel(_accelerator);
 		_snapshot = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot();
 	}
 
@@ -241,7 +243,7 @@ public sealed class DivisorCycleCache
 	[ThreadStatic]
 	private static MemoryBuffer1D<byte, Stride1D.Dense>? statusBuffer;
 
-	private static void EnsureCapacity(int requiredCapacity)
+	private void EnsureCapacity(int requiredCapacity)
 	{
 		if (requiredCapacity <= (divisorBuffer?.Length ?? 0))
 		{
@@ -257,11 +259,15 @@ public sealed class DivisorCycleCache
 			statusBuffer!.Dispose();			
 		}
 
-		divisorBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
-		powBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
-		orderBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
-		resultBuffer = SharedGpuContext.Accelerator.Allocate1D<ulong>(requiredCapacity);
-		statusBuffer = SharedGpuContext.Accelerator.Allocate1D<byte>(requiredCapacity);
+		Accelerator accelerator = _accelerator;
+		// lock(accelerator)
+		{
+			divisorBuffer = accelerator.Allocate1D<ulong>(requiredCapacity);
+			powBuffer = accelerator.Allocate1D<ulong>(requiredCapacity);
+			orderBuffer = accelerator.Allocate1D<ulong>(requiredCapacity);
+			resultBuffer = accelerator.Allocate1D<ulong>(requiredCapacity);
+			statusBuffer = accelerator.Allocate1D<byte>(requiredCapacity);			
+		}
 	}
 
 	private void ComputeCyclesGpuCore(ReadOnlySpan<ulong> divisors, Span<ulong> destination)
@@ -269,7 +275,7 @@ public sealed class DivisorCycleCache
 		int length = divisors.Length;
 		GpuPrimeWorkLimiter.Acquire();
 
-		Accelerator accelerator = SharedGpuContext.Accelerator;
+		Accelerator accelerator = _accelerator;
 		var stream = accelerator.CreateStream();
 		var kernel = _gpuKernel;
 		EnsureCapacity(length);

PerfectNumbers.Core/Gpu/AcceleratorPool.cs

@@ -0,0 +1,33 @@
+using ILGPU;
+using ILGPU.Runtime;
+
+namespace PerfectNumbers.Core.Gpu
+{
+    public struct AcceleratorPool(int capacity)
+	{
+		internal readonly Accelerator[] Accelerators = [.. Enumerable.Range(1, capacity).Select(_ => {
+			return SharedGpuContext.CreateAccelerator();
+		})];
+
+		private int _index;
+
+		public Accelerator Rent()
+		{
+			Atomic.CompareExchange(ref _index, capacity, 0);
+			var index = Atomic.Add(ref _index, 1);
+			if (index >= capacity)
+			{
+				index -= capacity;
+			}
+
+			return Accelerators[index];
+		}
+
+		public void Return(Accelerator accelerator)
+		{
+			// Intentionally left empty - there's nothing to do here. We're not really renting anything.
+		}
+
+		public static readonly AcceleratorPool Shared = new(PerfectNumberConstants.RollingAccelerators);
+	}
+}

PerfectNumbers.Core/Gpu/Accelerators/Pow2MontgomeryAccelerator.cs

@@ -0,0 +1,74 @@
+using System.Runtime.CompilerServices;
+using ILGPU;
+using ILGPU.Runtime;
+using PerfectNumbers.Core.Gpu.Kernels;
+
+namespace PerfectNumbers.Core.Gpu.Accelerators;
+
+public sealed class Pow2MontgomeryAccelerator
+{
+	#region Pool
+	
+	[ThreadStatic]
+	private static Queue<Pow2MontgomeryAccelerator>? _pool;
+
+	private static AcceleratorPool _accelerators = new(PerfectNumberConstants.RollingAccelerators);
+
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static Pow2MontgomeryAccelerator Rent()
+	{
+		var pool = _pool ??= [];
+		if (!pool.TryDequeue(out var gpu))
+		{
+			Accelerator accelerator = _accelerators.Rent();
+			gpu = new(accelerator);
+		}
+
+		return gpu;
+	}
+
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static void Return(Pow2MontgomeryAccelerator gpu)
+	{
+		gpu.Stream!.Dispose();
+		gpu.Stream = null;
+		_pool!.Enqueue(gpu);
+	}
+
+	#endregion
+
+	public readonly Accelerator Accelerator;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> Input;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> Output;
+	public AcceleratorStream? Stream;
+
+	public readonly Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> ConvertToStandardKernel;
+	public readonly Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> KeepMontgomeryKernel;
+
+	public Pow2MontgomeryAccelerator(Accelerator accelerator)
+	{
+		Accelerator = accelerator;
+		Input = accelerator.Allocate1D<ulong>(1);
+		Output = accelerator.Allocate1D<ulong>(1);
+		Stream = accelerator.CreateStream();
+
+		var keepKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2MontgomeryKernels.Pow2MontgomeryKernelKeepMontgomery);
+		var keepLauncher = KernelUtil.GetKernel(keepKernel);
+
+		KeepMontgomeryKernel = keepLauncher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
+
+		var convertKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2MontgomeryKernels.Pow2MontgomeryKernelConvertToStandard);
+		var convertLauncher = KernelUtil.GetKernel(convertKernel);
+
+		ConvertToStandardKernel = convertLauncher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
+	}
+
+	public void Dispose()
+	{
+		Input.Dispose();
+		Output.Dispose();
+		Stream?.Dispose();
+		// Accelerator is shared with other threads
+		// Accelerator.Dispose();
+	}
+}

PerfectNumbers.Core/Gpu/ExponentRemainderStepperGpu.cs

@@ -24,17 +24,14 @@ internal struct ExponentRemainderStepperGpu
         _currentResidue = GpuUInt128.One;
     }
 
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public bool MatchesDivisor(in GpuDivisorPartialData divisor)
-    {
-        return _modulus == divisor.Modulus;
-    }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public readonly bool MatchesDivisor(in GpuDivisorPartialData divisor) => _modulus == divisor.Modulus;
 
-    /// <summary>
-    /// Clears the cached residue so the stepper can be reused for another divisor.
-    /// Callers must reinitialize the stepper before consuming any residues after a reset.
-    /// </summary>
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+	/// <summary>
+	/// Clears the cached residue so the stepper can be reused for another divisor.
+	/// Callers must reinitialize the stepper before consuming any residues after a reset.
+	/// </summary>
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
     public void Reset()
     {
         PreviousExponent = 0UL;

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -5,7 +5,8 @@ namespace PerfectNumbers.Core.Gpu;
 
 public class GpuKernelPool
 {
-	public static readonly KernelContainer Kernels = CreateKernels();
+	[ThreadStatic]
+	private static Dictionary<Accelerator, KernelContainer>? _pool;
 
 	public static T InitOnce<T>(ref T? target, Func<T> valueFactory) where T : class
 	{
@@ -18,14 +19,19 @@ public class GpuKernelPool
 		return Interlocked.CompareExchange(ref target, newValue, null) ?? newValue;
 	}
 
-	private static KernelContainer CreateKernels()
+	internal static KernelContainer GetOrAddKernels(Accelerator accelerator, AcceleratorStream stream)
 	{
+		var pool = _pool ??= [];
+		if (pool.TryGetValue(accelerator, out var kernelContainer))
+		{
+			return kernelContainer;
+		}
+
 		var kernels = new KernelContainer();
-		var stream = SharedGpuContext.Accelerator.CreateStream();
 
 		InitOnce(ref kernels.Order, () =>
 		{
-			var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>(OrderKernels.OrderKernelScan);
+			var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>(OrderKernels.OrderKernelScan);
 
 			var kernel = KernelUtil.GetKernel(loaded);
 
@@ -34,7 +40,7 @@ public class GpuKernelPool
 
 		InitOnce(ref kernels.Pow2Mod, () =>
 		{
-			var loaded = SharedGpuContext.Accelerator.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 loaded = accelerator.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);
 
@@ -43,7 +49,7 @@ public class GpuKernelPool
 
 		InitOnce(ref kernels.Incremental, () =>
 		{
-			var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalKernelScan);
+			var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalKernelScan);
 
 			var kernel = KernelUtil.GetKernel(loaded);
 
@@ -52,7 +58,7 @@ public class GpuKernelPool
 
 		InitOnce(ref kernels.IncrementalOrder, () =>
 		{
-			var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalOrderKernelScan);
+			var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalOrderKernelScan);
 
 			var kernel = KernelUtil.GetKernel(loaded);
 
@@ -62,7 +68,7 @@ public class GpuKernelPool
 
 		InitOnce(ref kernels.Pow2ModOrder, () =>
 		{
-			var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernels.Pow2ModOrderKernelScan);
+			var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernels.Pow2ModOrderKernelScan);
 
 			var kernel = KernelUtil.GetKernel(loaded);
 
@@ -71,7 +77,7 @@ public class GpuKernelPool
 
 		InitOnce(ref kernels.SpecialMax, () =>
 		{
-			var loaded = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel);
+			var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel);
 
 			var kernel = KernelUtil.GetKernel(loaded);
 
@@ -80,126 +86,132 @@ public class GpuKernelPool
 		
 		InitOnce(ref kernels.SmallPrimeFactor, () =>
 		{
-			var loaded = SharedGpuContext.Accelerator.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 loaded = accelerator.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);
 			
 			return kernel.CreateLauncherDelegate<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>>>();
 		});
 			
-		PreloadStaticTables(kernels, stream);
+		PreloadStaticTables(accelerator, kernels, stream);
 
-		stream.Synchronize();
-		stream.Dispose();
+		pool[accelerator] = kernels;
 		return kernels;
 	}
 
-	private static KernelContainer CreateKernels(AcceleratorStream stream)
+	internal static void PreloadStaticTables(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream)
 	{
-		KernelContainer kernels = new();
-		PreloadStaticTables(kernels, stream);
-		return Kernels;
+		EnsureSmallCyclesOnDevice(accelerator, kernels, stream);
+		EnsureSmallPrimesOnDevice(accelerator, kernels, stream);
+		EnsureSmallPrimeFactorTables(accelerator, kernels, stream);
 	}
 
-	internal static void PreloadStaticTables(KernelContainer kernels, AcceleratorStream stream)
-	{
-		_ = EnsureSmallCyclesOnDevice(kernels, stream);
-		_ = EnsureSmallPrimesOnDevice(kernels, stream);
-		_ = EnsureSmallPrimeFactorTables(kernels, stream);
-	}
-
-	public static ArrayView1D<ulong, Stride1D.Dense> GetSmallCyclesOnDevice() => Kernels.SmallCycles!;
+	public static ArrayView1D<ulong, Stride1D.Dense> GetSmallCyclesOnDevice(KernelContainer kernels) => kernels.SmallCycles!;
 
 	// Ensures the small cycles table is uploaded to the device for the given accelerator.
 	// Returns the ArrayView to pass into kernels (when kernels are extended to accept it).
-	private static ArrayView1D<ulong, Stride1D.Dense> EnsureSmallCyclesOnDevice(KernelContainer kernels, AcceleratorStream stream)
+	private static void EnsureSmallCyclesOnDevice(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream)
 	{
-		if (kernels.SmallCycles is { } buffer)
+		if (kernels.SmallCycles is { })
 		{
-			return buffer.View;
+			return;
+		}
+
+		var host = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot();
+		MemoryBuffer1D<ulong, Stride1D.Dense>? device;
+
+		// lock (accelerator)
+		{
+			device = accelerator.Allocate1D<ulong>(host.Length);
 		}
 
-		var host = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot(); // TODO: Preload this device buffer during startup and keep it immutable so we can delete the lock above in favor of the preloaded snapshot.
-		var device = SharedGpuContext.Accelerator.Allocate1D<ulong>(host.Length);
 		device.View.CopyFromCPU(stream, host);
 		kernels.SmallCycles = device;
-		return device.View;
 	}
 
-	public static ResiduePrimeViews GetSmallPrimesOnDevice() => new(
-		Kernels.SmallPrimesLastOne!,
-		Kernels.SmallPrimesLastSeven!,
-		Kernels.SmallPrimesPow2LastOne!,
-		Kernels.SmallPrimesPow2LastSeven!
+	public static ResiduePrimeViews GetSmallPrimesOnDevice(KernelContainer kernels) => new(
+		kernels.SmallPrimesLastOne!,
+		kernels.SmallPrimesLastSeven!,
+		kernels.SmallPrimesPow2LastOne!,
+		kernels.SmallPrimesPow2LastSeven!
 	);
 
-	private static ResiduePrimeViews EnsureSmallPrimesOnDevice(KernelContainer kernels, AcceleratorStream stream)
+	private static void EnsureSmallPrimesOnDevice(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream)
 	{
-		if (kernels.SmallPrimesLastOne is { } lastOne &&
-			kernels.SmallPrimesLastSeven is { } lastSeven &&
-			kernels.SmallPrimesPow2LastOne is { } lastOnePow2 &&
-			kernels.SmallPrimesPow2LastSeven is { } lastSevenPow2)
+		if (kernels.SmallPrimesLastOne is { } &&
+			kernels.SmallPrimesLastSeven is { } &&
+			kernels.SmallPrimesPow2LastOne is { } &&
+			kernels.SmallPrimesPow2LastSeven is { })
 		{
-			return new ResiduePrimeViews(lastOne.View, lastSeven.View, lastOnePow2.View, lastSevenPow2.View);
+			return;
 		}
 
-		if (kernels.SmallPrimesLastOne is { } existingLastOne &&
-			kernels.SmallPrimesLastSeven is { } existingLastSeven &&
-			kernels.SmallPrimesPow2LastOne is { } existingLastOnePow2 &&
-			kernels.SmallPrimesPow2LastSeven is { } existingLastSevenPow2)
+		if (kernels.SmallPrimesLastOne is { } &&
+			kernels.SmallPrimesLastSeven is { } &&
+			kernels.SmallPrimesPow2LastOne is { } &&
+			kernels.SmallPrimesPow2LastSeven is { })
 		{
-			return new ResiduePrimeViews(existingLastOne.View, existingLastSeven.View, existingLastOnePow2.View, existingLastSevenPow2.View);
+			return;
 		}
 
-		Accelerator accelerator = SharedGpuContext.Accelerator;
 		var hostLastOne = PrimesGenerator.SmallPrimesLastOne;
 		var hostLastSeven = PrimesGenerator.SmallPrimesLastSeven;
 		var hostLastOnePow2 = PrimesGenerator.SmallPrimesPow2LastOne;
 		var hostLastSevenPow2 = PrimesGenerator.SmallPrimesPow2LastSeven;
 
-		var deviceLastOne = accelerator.Allocate1D<uint>(hostLastOne.Length);
-		deviceLastOne.View.CopyFromCPU(stream, hostLastOne);
-		var deviceLastSeven = accelerator.Allocate1D<uint>(hostLastSeven.Length);
-		deviceLastSeven.View.CopyFromCPU(stream, hostLastSeven);
-		var deviceLastOnePow2 = accelerator.Allocate1D<ulong>(hostLastOnePow2.Length);
-		deviceLastOnePow2.View.CopyFromCPU(stream, hostLastOnePow2);
-		var deviceLastSevenPow2 = accelerator.Allocate1D<ulong>(hostLastSevenPow2.Length);
-		deviceLastSevenPow2.View.CopyFromCPU(stream, hostLastSevenPow2);
+		MemoryBuffer1D<uint, Stride1D.Dense>? deviceLastOne;
+		MemoryBuffer1D<uint, Stride1D.Dense>? deviceLastSeven;
+		MemoryBuffer1D<ulong, Stride1D.Dense>? deviceLastOnePow2;
+		MemoryBuffer1D<ulong, Stride1D.Dense>? deviceLastSevenPow2;
 
+		// lock (accelerator)
+		// {
+			deviceLastOne = accelerator.Allocate1D<uint>(hostLastOne.Length);
+			deviceLastSeven = accelerator.Allocate1D<uint>(hostLastSeven.Length);
+			deviceLastOnePow2 = accelerator.Allocate1D<ulong>(hostLastOnePow2.Length);
+			deviceLastSevenPow2 = accelerator.Allocate1D<ulong>(hostLastSevenPow2.Length);
+		// }
+		
+		deviceLastOne.View.CopyFromCPU(stream, hostLastOne);
+		deviceLastSeven.View.CopyFromCPU(stream, hostLastSeven);
+		deviceLastOnePow2.View.CopyFromCPU(stream, hostLastOnePow2);
+		deviceLastSevenPow2.View.CopyFromCPU(stream, hostLastSevenPow2);
 		kernels.SmallPrimesLastOne = deviceLastOne;
 		kernels.SmallPrimesLastSeven = deviceLastSeven;
 		kernels.SmallPrimesPow2LastOne = deviceLastOnePow2;
 		kernels.SmallPrimesPow2LastSeven = deviceLastSevenPow2;
-
-		return new ResiduePrimeViews(deviceLastOne.View, deviceLastSeven.View, deviceLastOnePow2.View, deviceLastSevenPow2.View);
 	}
 
-	public static SmallPrimeFactorViews GetSmallPrimeFactorTables() => new(
-		Kernels.SmallPrimeFactorsPrimes!,
-		Kernels.SmallPrimeFactorsSquares!
+	public static SmallPrimeFactorViews GetSmallPrimeFactorTables(KernelContainer kernels) => new(
+		kernels.SmallPrimeFactorsPrimes!,
+		kernels.SmallPrimeFactorsSquares!
 	);
 	
-	private static SmallPrimeFactorViews EnsureSmallPrimeFactorTables(KernelContainer kernels, AcceleratorStream stream)
+	private static void EnsureSmallPrimeFactorTables(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream)
 	{
-		if (kernels.SmallPrimeFactorsPrimes is { } primeBuffer && kernels.SmallPrimeFactorsSquares is { } squareBuffer)
+		if (kernels.SmallPrimeFactorsPrimes is { } && kernels.SmallPrimeFactorsSquares is { })
 		{
-			return new SmallPrimeFactorViews(primeBuffer, squareBuffer);
+			return;
 		}
 
-		Accelerator accelerator = SharedGpuContext.Accelerator;
 		var hostPrimes = PrimesGenerator.SmallPrimes;
 		var hostSquares = PrimesGenerator.SmallPrimesPow2;
 
 		int length = hostPrimes.Length;
-		var devicePrimes = accelerator.Allocate1D<uint>(length);
+		MemoryBuffer1D<uint, Stride1D.Dense>? devicePrimes;
+		MemoryBuffer1D<ulong, Stride1D.Dense>? deviceSquares;
+
+		// lock (accelerator)
+		{
+			devicePrimes = accelerator.Allocate1D<uint>(length);
+			deviceSquares = accelerator.Allocate1D<ulong>(length);
+		}
+		
 		devicePrimes.View.CopyFromCPU(stream, hostPrimes);
-		var deviceSquares = accelerator.Allocate1D<ulong>(length);
 		deviceSquares.View.CopyFromCPU(stream, hostSquares);
 
 		kernels.SmallPrimeFactorsPrimes = devicePrimes;
 		kernels.SmallPrimeFactorsSquares = deviceSquares;
-
-		return new SmallPrimeFactorViews(devicePrimes, deviceSquares);
 	}
 
 	/// <summary>
@@ -209,11 +221,12 @@ public class GpuKernelPool
 	public static void Run(Action<Accelerator, AcceleratorStream> action)
 	{
 		GpuPrimeWorkLimiter.Acquire();
-		var accelerator = SharedGpuContext.Accelerator;
+		var accelerator = SharedGpuContext.CreateAccelerator();
 		var stream = accelerator.CreateStream();
 		action(accelerator, stream);
 		stream.Synchronize();
 		stream.Dispose();
+		accelerator.Dispose();
 		GpuPrimeWorkLimiter.Release();
 	}
 }

PerfectNumbers.Core/Gpu/GpuScratchBufferPool.cs

@@ -22,7 +22,7 @@ public sealed class ScratchBuffer
 		_accelerator = accelerator;
 
 		// lock (accelerator)
-		// {
+		{
 			SmallPrimeFactorPrimeSlots = accelerator.Allocate1D<ulong>(smallPrimeFactorSlotCount);
 			SmallPrimeFactorExponentSlots = accelerator.Allocate1D<int>(smallPrimeFactorSlotCount);
 			SmallPrimeFactorCountSlot = accelerator.Allocate1D<int>(1);
@@ -31,7 +31,7 @@ public sealed class ScratchBuffer
 			SpecialMaxFactors = accelerator.Allocate1D<ulong>(specialMaxFactorCapacity);
 			SpecialMaxCandidates = accelerator.Allocate1D<ulong>(specialMaxFactorCapacity);
 			SpecialMaxResult = accelerator.Allocate1D<ulong>(1);
-		// }
+		}
 	}
 
 	public void ResizeSmallPrimeFactorSlots(int newSize)
@@ -41,8 +41,11 @@ public sealed class ScratchBuffer
 			SmallPrimeFactorPrimeSlots.Dispose();
 			SmallPrimeFactorExponentSlots.Dispose();
 
-			SmallPrimeFactorPrimeSlots = _accelerator.Allocate1D<ulong>(newSize);
-			SmallPrimeFactorExponentSlots = _accelerator.Allocate1D<int>(newSize);
+			// lock(_accelerator)
+			{
+				SmallPrimeFactorPrimeSlots = _accelerator.Allocate1D<ulong>(newSize);
+				SmallPrimeFactorExponentSlots = _accelerator.Allocate1D<int>(newSize);				
+			}
 		}
 	}
 
@@ -53,8 +56,11 @@ public sealed class ScratchBuffer
 			SpecialMaxFactors.Dispose();
 			SpecialMaxCandidates.Dispose();
 
-			SpecialMaxFactors = _accelerator.Allocate1D<ulong>(newSize);
-			SpecialMaxCandidates = _accelerator.Allocate1D<ulong>(newSize);
+			// lock(_accelerator)
+			{
+				SpecialMaxFactors = _accelerator.Allocate1D<ulong>(newSize);
+				SpecialMaxCandidates = _accelerator.Allocate1D<ulong>(newSize);				
+			}
 		}
 	}
 }
@@ -63,6 +69,7 @@ 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
@@ -89,13 +96,13 @@ public static class GpuScratchBufferPool
 		{
 			if (buffer.SmallPrimeFactorPrimeSlots.Length < smallPrimeFactorSlotCount || buffer.SpecialMaxFactors.Length < specialMaxFactorCapacity)
 			{
-				throw new InvalidOperationException();
+				// 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);
-				// }
+					buffer.ResizeSmallPrimeFactorSlots(smallPrimeFactorSlotCount);
+					buffer.ResizeSpecialMaxFactors(specialMaxFactorCapacity);
+				}
 			}
 			return buffer;
 		}
@@ -113,12 +120,12 @@ public static class GpuScratchBufferPool
 
 	public static void WarmUp(int count, int smallPrimeFactorSlots, int specialMaxFactorCapacity)
 	{
-		Accelerator accelerator = SharedGpuContext.Accelerator;
+		// Accelerator accelerator = SharedGpuContext.Accelerator;
 
-		for (int i = 0; i < count; i++)
-		{
-			var buffer = Rent(accelerator, smallPrimeFactorSlots, specialMaxFactorCapacity);
-			Return(buffer);
-		}
+		// for (int i = 0; i < count; i++)
+		// {
+		// 	var buffer = Rent(accelerator, smallPrimeFactorSlots, specialMaxFactorCapacity);
+		// 	Return(buffer);
+		// }
 	}
 }

PerfectNumbers.Core/Gpu/GpuStaticTableInitializer.cs

@@ -1,25 +1,15 @@
-using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Runtime.CompilerServices;
-using ILGPU;
 using ILGPU.Runtime;
-using PerfectNumbers.Core;
 
 namespace PerfectNumbers.Core.Gpu;
 
 internal static class GpuStaticTableInitializer
 {
-    private static readonly ConcurrentDictionary<Accelerator, byte> Initialized = new(AcceleratorReferenceComparer.Instance);
-
-    internal static void EnsureStaticTables(KernelContainer kernels, Accelerator accelerator, AcceleratorStream stream)
-    {
-        Initialized.GetOrAdd(accelerator, acc =>
-        {
-            PrimeTester.PrimeTesterGpuContextPool.EnsureStaticTables(kernels, acc, stream);
-            GpuKernelPool.PreloadStaticTables(kernels, stream);
-            PrimeOrderGpuHeuristics.PreloadStaticTables( acc, stream);
-            return 0;
-        });
+    internal static void EnsureStaticTables(Accelerator accelerator, KernelContainer kernels, AcceleratorStream stream)
+	{
+		// PrimeTester.PrimeTesterGpuContextPool.EnsureStaticTables(accelerator);
+		GpuKernelPool.PreloadStaticTables(accelerator, kernels, stream);
+		// PrimeOrderGpuHeuristics.PreloadStaticTables(accelerator, stream);
     }
 
     private sealed class AcceleratorReferenceComparer : IEqualityComparer<Accelerator>

PerfectNumbers.Core/Gpu/HeuristicGpuDivisorTables.cs

@@ -6,141 +6,134 @@ namespace PerfectNumbers.Core.Gpu;
 
 internal readonly struct HeuristicGpuDivisorTables
 {
-        public static readonly HeuristicGpuDivisorTables Shared = default;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding1;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding3;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding7;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding9;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding1;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding3;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding7;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding9;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupADivisors;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupADivisorSquares;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorsEnding1;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorSquaresEnding1;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorsEnding7;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorSquaresEnding7;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorsEnding9;
+	public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorSquaresEnding9;
 
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding1;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding3;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding7;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorsEnding9;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding1;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding3;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding7;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> CombinedDivisorSquaresEnding9;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupADivisors;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupADivisorSquares;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorsEnding1;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorSquaresEnding1;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorsEnding7;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorSquaresEnding7;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorsEnding9;
-        public readonly ArrayView1D<ulong, Stride1D.Dense> GroupBDivisorSquaresEnding9;
+	public HeuristicGpuDivisorTables(
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding1,
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding3,
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding7,
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding9,
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding1,
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding3,
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding7,
+			ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding9,
+			ArrayView1D<ulong, Stride1D.Dense> groupADivisors,
+			ArrayView1D<ulong, Stride1D.Dense> groupADivisorSquares,
+			ArrayView1D<ulong, Stride1D.Dense> groupBDivisorsEnding1,
+			ArrayView1D<ulong, Stride1D.Dense> groupBDivisorSquaresEnding1,
+			ArrayView1D<ulong, Stride1D.Dense> groupBDivisorsEnding7,
+			ArrayView1D<ulong, Stride1D.Dense> groupBDivisorSquaresEnding7,
+			ArrayView1D<ulong, Stride1D.Dense> groupBDivisorsEnding9,
+			ArrayView1D<ulong, Stride1D.Dense> groupBDivisorSquaresEnding9)
+	{
+		CombinedDivisorsEnding1 = combinedDivisorsEnding1;
+		CombinedDivisorsEnding3 = combinedDivisorsEnding3;
+		CombinedDivisorsEnding7 = combinedDivisorsEnding7;
+		CombinedDivisorsEnding9 = combinedDivisorsEnding9;
+		CombinedDivisorSquaresEnding1 = combinedDivisorSquaresEnding1;
+		CombinedDivisorSquaresEnding3 = combinedDivisorSquaresEnding3;
+		CombinedDivisorSquaresEnding7 = combinedDivisorSquaresEnding7;
+		CombinedDivisorSquaresEnding9 = combinedDivisorSquaresEnding9;
+		GroupADivisors = groupADivisors;
+		GroupADivisorSquares = groupADivisorSquares;
+		GroupBDivisorsEnding1 = groupBDivisorsEnding1;
+		GroupBDivisorSquaresEnding1 = groupBDivisorSquaresEnding1;
+		GroupBDivisorsEnding7 = groupBDivisorsEnding7;
+		GroupBDivisorSquaresEnding7 = groupBDivisorSquaresEnding7;
+		GroupBDivisorsEnding9 = groupBDivisorsEnding9;
+		GroupBDivisorSquaresEnding9 = groupBDivisorSquaresEnding9;
+	}
 
-        public HeuristicGpuDivisorTables(
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding1,
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding3,
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding7,
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorsEnding9,
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding1,
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding3,
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding7,
-                ArrayView1D<ulong, Stride1D.Dense> combinedDivisorSquaresEnding9,
-                ArrayView1D<ulong, Stride1D.Dense> groupADivisors,
-                ArrayView1D<ulong, Stride1D.Dense> groupADivisorSquares,
-                ArrayView1D<ulong, Stride1D.Dense> groupBDivisorsEnding1,
-                ArrayView1D<ulong, Stride1D.Dense> groupBDivisorSquaresEnding1,
-                ArrayView1D<ulong, Stride1D.Dense> groupBDivisorsEnding7,
-                ArrayView1D<ulong, Stride1D.Dense> groupBDivisorSquaresEnding7,
-                ArrayView1D<ulong, Stride1D.Dense> groupBDivisorsEnding9,
-                ArrayView1D<ulong, Stride1D.Dense> groupBDivisorSquaresEnding9)
-        {
-                CombinedDivisorsEnding1 = combinedDivisorsEnding1;
-                CombinedDivisorsEnding3 = combinedDivisorsEnding3;
-                CombinedDivisorsEnding7 = combinedDivisorsEnding7;
-                CombinedDivisorsEnding9 = combinedDivisorsEnding9;
-                CombinedDivisorSquaresEnding1 = combinedDivisorSquaresEnding1;
-                CombinedDivisorSquaresEnding3 = combinedDivisorSquaresEnding3;
-                CombinedDivisorSquaresEnding7 = combinedDivisorSquaresEnding7;
-                CombinedDivisorSquaresEnding9 = combinedDivisorSquaresEnding9;
-                GroupADivisors = groupADivisors;
-                GroupADivisorSquares = groupADivisorSquares;
-                GroupBDivisorsEnding1 = groupBDivisorsEnding1;
-                GroupBDivisorSquaresEnding1 = groupBDivisorSquaresEnding1;
-                GroupBDivisorsEnding7 = groupBDivisorsEnding7;
-                GroupBDivisorSquaresEnding7 = groupBDivisorSquaresEnding7;
-                GroupBDivisorsEnding9 = groupBDivisorsEnding9;
-                GroupBDivisorSquaresEnding9 = groupBDivisorSquaresEnding9;
-        }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public ArrayView1D<ulong, Stride1D.Dense> SelectCombined(byte ending)
+	{
+		return ending switch
+		{
+			1 => CombinedDivisorsEnding1,
+			3 => CombinedDivisorsEnding3,
+			7 => CombinedDivisorsEnding7,
+			9 => CombinedDivisorsEnding9,
+			_ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
+		};
+	}
 
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public ArrayView1D<ulong, Stride1D.Dense> SelectCombined(byte ending)
-        {
-                return ending switch
-                {
-                        1 => CombinedDivisorsEnding1,
-                        3 => CombinedDivisorsEnding3,
-                        7 => CombinedDivisorsEnding7,
-                        9 => CombinedDivisorsEnding9,
-                        _ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
-                };
-        }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public ArrayView1D<ulong, Stride1D.Dense> SelectCombinedSquares(byte ending)
+	{
+		return ending switch
+		{
+			1 => CombinedDivisorSquaresEnding1,
+			3 => CombinedDivisorSquaresEnding3,
+			7 => CombinedDivisorSquaresEnding7,
+			9 => CombinedDivisorSquaresEnding9,
+			_ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
+		};
+	}
 
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public ArrayView1D<ulong, Stride1D.Dense> SelectCombinedSquares(byte ending)
-        {
-                return ending switch
-                {
-                        1 => CombinedDivisorSquaresEnding1,
-                        3 => CombinedDivisorSquaresEnding3,
-                        7 => CombinedDivisorSquaresEnding7,
-                        9 => CombinedDivisorSquaresEnding9,
-                        _ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
-                };
-        }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public ArrayView1D<ulong, Stride1D.Dense> SelectGroupB(byte ending)
+	{
+		return ending switch
+		{
+			1 => GroupBDivisorsEnding1,
+			7 => GroupBDivisorsEnding7,
+			9 => GroupBDivisorsEnding9,
+			_ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
+		};
+	}
 
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public ArrayView1D<ulong, Stride1D.Dense> SelectGroupB(byte ending)
-        {
-                return ending switch
-                {
-                        1 => GroupBDivisorsEnding1,
-                        7 => GroupBDivisorsEnding7,
-                        9 => GroupBDivisorsEnding9,
-                        _ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
-                };
-        }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public ArrayView1D<ulong, Stride1D.Dense> SelectGroupBSquares(byte ending)
+	{
+		return ending switch
+		{
+			1 => GroupBDivisorSquaresEnding1,
+			7 => GroupBDivisorSquaresEnding7,
+			9 => GroupBDivisorSquaresEnding9,
+			_ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
+		};
+	}
 
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public ArrayView1D<ulong, Stride1D.Dense> SelectGroupBSquares(byte ending)
-        {
-                return ending switch
-                {
-                        1 => GroupBDivisorSquaresEnding1,
-                        7 => GroupBDivisorSquaresEnding7,
-                        9 => GroupBDivisorSquaresEnding9,
-                        _ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
-                };
-        }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public ArrayView1D<ulong, Stride1D.Dense> SelectDivisors(HeuristicGpuDivisorTableKind kind, byte ending)
+	{
+		return kind switch
+		{
+			HeuristicGpuDivisorTableKind.GroupA => GroupADivisors,
+			HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorsEnding1,
+			HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorsEnding7,
+			HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorsEnding9,
+			HeuristicGpuDivisorTableKind.Combined => SelectCombined(ending),
+			_ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
+		};
+	}
 
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public ArrayView1D<ulong, Stride1D.Dense> SelectDivisors(HeuristicGpuDivisorTableKind kind, byte ending)
-        {
-                return kind switch
-                {
-                        HeuristicGpuDivisorTableKind.GroupA => GroupADivisors,
-                        HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorsEnding1,
-                        HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorsEnding7,
-                        HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorsEnding9,
-                        HeuristicGpuDivisorTableKind.Combined => SelectCombined(ending),
-                        _ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
-                };
-        }
-
-        [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public ArrayView1D<ulong, Stride1D.Dense> SelectDivisorSquares(HeuristicGpuDivisorTableKind kind, byte ending)
-        {
-                return kind switch
-                {
-                        HeuristicGpuDivisorTableKind.GroupA => GroupADivisorSquares,
-                        HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorSquaresEnding1,
-                        HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorSquaresEnding7,
-                        HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorSquaresEnding9,
-                        HeuristicGpuDivisorTableKind.Combined => SelectCombinedSquares(ending),
-                        _ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
-                };
-        }
-
-        internal static void InitializeShared(in HeuristicGpuDivisorTables tables)
-        {
-                Unsafe.AsRef(in Shared) = tables;
-        }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public ArrayView1D<ulong, Stride1D.Dense> SelectDivisorSquares(HeuristicGpuDivisorTableKind kind, byte ending)
+	{
+		return kind switch
+		{
+			HeuristicGpuDivisorTableKind.GroupA => GroupADivisorSquares,
+			HeuristicGpuDivisorTableKind.GroupBEnding1 => GroupBDivisorSquaresEnding1,
+			HeuristicGpuDivisorTableKind.GroupBEnding7 => GroupBDivisorSquaresEnding7,
+			HeuristicGpuDivisorTableKind.GroupBEnding9 => GroupBDivisorSquaresEnding9,
+			HeuristicGpuDivisorTableKind.Combined => SelectCombinedSquares(ending),
+			_ => ArrayView1D<ulong, Stride1D.Dense>.Empty,
+		};
+	}
 }

PerfectNumbers.Core/Gpu/InputOutputScratchBuffers.cs

@@ -0,0 +1,126 @@
+using System.Collections.Concurrent;
+using ILGPU;
+using ILGPU.Runtime;
+
+namespace PerfectNumbers.Core.Gpu;
+
+public sealed class GpuInputOutputScratchBuffer<TInput, TOutput>
+	where TInput : unmanaged
+	where TOutput: unmanaged
+{
+	private readonly Accelerator _accelerator;
+	public MemoryBuffer1D<TInput, Stride1D.Dense> Input { get; private set; }
+	public MemoryBuffer1D<TOutput, Stride1D.Dense> Output { get; private set; }
+
+	public GpuInputOutputScratchBuffer(Accelerator accelerator,int inputCapacity, int outputCapacity)
+	{
+		_accelerator = accelerator;
+		// lock (accelerator)
+		{
+			Input = accelerator.Allocate1D<TInput>(inputCapacity);
+			Output = accelerator.Allocate1D<TOutput>(outputCapacity);
+		}
+	}
+
+	public void EnsureCapacity(int inputCapacity, int outputCapacity)
+	{
+		if (Input.Length >= inputCapacity && Output.Length >= outputCapacity)
+		{
+			return;
+		}
+
+		Accelerator accelerator = _accelerator;
+		if (Input.Length >= inputCapacity && Output.Length >= outputCapacity)
+		{
+			return;
+		}
+
+		// lock (accelerator)
+		{
+			if (Input.Length < inputCapacity)
+			{
+				Input.Dispose();
+				Input = accelerator.Allocate1D<TInput>(inputCapacity);
+			}
+
+			if (Output.Length < outputCapacity)
+			{
+				Output.Dispose();
+				Output = accelerator.Allocate1D<TOutput>(outputCapacity);
+			}
+		}
+	}
+	
+	public void Dispose()
+	{
+		Input.Dispose();
+		Output.Dispose();
+	}
+}
+
+public static class GpuInputOutputScratchBufferPool<TInput, TOutput>
+	where TInput : unmanaged
+	where TOutput : unmanaged
+{
+	private const int WarmUpCapacity = 20480;
+
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Queue<GpuInputOutputScratchBuffer<TInput, TOutput>>>? _pool;
+
+	private static Queue<GpuInputOutputScratchBuffer<TInput, TOutput>> GetPools(Accelerator accelerator)
+	{
+		var pool = _pool ??= [];
+
+		if (!pool.TryGetValue(accelerator, out var bufferQueue))
+		{
+			bufferQueue = [];
+			pool[accelerator] = bufferQueue;
+		}
+		
+		return bufferQueue;
+	}
+
+	// private static readonly ConcurrentDictionary<Accelerator, ConcurrentQueue<InputOutputScratchBuffer>> _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 GpuInputOutputScratchBuffer<TInput, TOutput> Rent(Accelerator accelerator, int inputCapacity, int outputCapacity)
+	{
+		var bufferQueue = GetPools(accelerator);
+		if (bufferQueue.TryDequeue(out var buffer))
+		{
+			if (buffer.Input.Length < inputCapacity || buffer.Output.Length < outputCapacity)
+			{
+				// lock (accelerator)
+				// {
+				// Console.WriteLine($"Resizing GPU scratch buffer from pool ({buffer.SmallPrimeFactorPrimeSlots.Length} / {smallPrimeFactorSlotCount}), ({buffer.SpecialMaxFactors.Length}/{specialMaxFactorCapacity})");
+				buffer.EnsureCapacity(inputCapacity, outputCapacity);
+				// }
+			}
+
+			return buffer;
+		}
+
+		return new GpuInputOutputScratchBuffer<TInput, TOutput>(accelerator, inputCapacity, outputCapacity);
+	}
+
+	public static void Return(Accelerator accelerator, GpuInputOutputScratchBuffer<TInput, TOutput> buffer) => GetPools(accelerator).Enqueue(buffer);
+
+	public static void WarmUp(int count, int inputCapacity, int outputCapacity)
+	{
+		// List<GpuInputOutputScratchBuffer<TInput, TOutput>> rented = new(count);
+		// for (int i = 0; i < count; i++)
+		// {
+		// 	rented.Add(Rent(inputCapacity, outputCapacity));
+		// }
+
+		// for (int i = 0; i < count; i++)
+		// {
+		// 	Return(rented[i]);
+		// }
+
+		// rented.Clear();
+		// rented.Capacity = 1;
+	}
+}

PerfectNumbers.Core/Gpu/Kernels/NttMontgomeryKernels.cs

@@ -30,7 +30,7 @@ internal static class NttMontgomeryKernels
         data[i2] = new GpuUInt128(0UL, diff);
     }
 
-    public static void ToMont64Kernel(Index1D index, AcceleratorStream stream, ArrayView<GpuUInt128> data, ulong modulus, ulong nPrime, ulong r2)
+    public static void ToMont64Kernel(Index1D index, ArrayView<GpuUInt128> data, ulong modulus, ulong nPrime, ulong r2)
     {
         ulong a = data[index].Low;
         ulong v = MontMul64(a, r2, modulus, nPrime);

PerfectNumbers.Core/Gpu/Kernels/Pow2MontgomeryKernels.cs

@@ -1,8 +1,7 @@
 using ILGPU;
 using ILGPU.Runtime;
-using PerfectNumbers.Core;
 
-namespace PerfectNumbers.Core.Gpu;
+namespace PerfectNumbers.Core.Gpu.Kernels;
 
 internal static class Pow2MontgomeryKernels
 {

PerfectNumbers.Core/Gpu/Kernels/PrimeOrderKernels.cs

@@ -1,10 +1,7 @@
-using System;
 using System.Numerics;
 using System.Runtime.CompilerServices;
 using ILGPU;
-using ILGPU.Algorithms.Random;
 using ILGPU.Runtime;
-using PerfectNumbers.Core;
 
 namespace PerfectNumbers.Core.Gpu;
 
@@ -1171,13 +1168,20 @@ internal static partial class PrimeOrderGpuHeuristics
 	}
 
     private static Action<AcceleratorStream, Index1D, ArrayView1D<GpuUInt128, Stride1D.Dense>, GpuUInt128, ArrayView1D<GpuUInt128, Stride1D.Dense>> GetPow2ModWideKernel(Accelerator accelerator)
-    {
-        return Pow2ModKernelWideCache.GetOrAdd(accelerator, static accel =>
-        {
-            var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<GpuUInt128, Stride1D.Dense>, GpuUInt128, ArrayView1D<GpuUInt128, Stride1D.Dense>>(Pow2ModKernelWide);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<GpuUInt128, Stride1D.Dense>, GpuUInt128, ArrayView1D<GpuUInt128, Stride1D.Dense>>>();
-        });
+	{
+		var pool = _pow2ModWideKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<GpuUInt128, Stride1D.Dense>, GpuUInt128, ArrayView1D<GpuUInt128, Stride1D.Dense>>(Pow2ModKernelWide);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<GpuUInt128, Stride1D.Dense>, GpuUInt128, ArrayView1D<GpuUInt128, Stride1D.Dense>>>();
+		pool[accelerator] = cached;
+		return cached;
     }
 
 	/// This kernel always sets the result of the corresponding element. Callers don't need to clear the output buffers.

PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs

@@ -22,22 +22,30 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
     // so the synchronization fields from the previous implementation remain commented out here.
     // private readonly object _sync = new();
     private ulong _divisorLimit;
-    // private bool _isConfigured;
+	// private bool _isConfigured;
 
-    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, 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<DivisorScanSession> _sessionPool = new();
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>>>? _kernel;
+	
+    private readonly ConcurrentBag<BatchResources> _resourcePool = [];
+    private readonly ConcurrentBag<DivisorScanSession> _sessionPool = [];
 
-	private Action<AcceleratorStream, Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>> GetKernel(Accelerator accelerator) =>
-		_kernelCache.GetOrAdd(accelerator, acc =>
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>> GetKernel(Accelerator accelerator)
+	{
+		var pool = _kernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>>(DivisorByDivisorKernels.CheckKernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>>>();
-		});
+			return cached;
+		}
 
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>>(DivisorByDivisorKernels.CheckKernel);
+		var kernel = KernelUtil.GetKernel(loaded);
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuDivisorPartialData>, ArrayView<int>, ArrayView<int>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>, ArrayView<int>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
 
-    public int GpuBatchSize
+	public int GpuBatchSize
     {
         get => _gpuBatchSize;
         set => _gpuBatchSize = value < 1 ? 1 : value;
@@ -96,7 +104,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         ulong lastProcessed;
 
         GpuPrimeWorkLimiter.Acquire();
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
 
         // Monitor.Enter(gpuLease.ExecutionLock);
@@ -502,7 +510,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         {
             _owner = owner;
             GpuPrimeWorkLimiter.Acquire();
-			Accelerator accelerator = SharedGpuContext.Accelerator;
+			Accelerator accelerator = AcceleratorPool.Shared.Rent();
             _accelerator = accelerator;
             _stream = accelerator.CreateStream();
             _capacity = Math.Max(1, owner._gpuBatchSize);
@@ -514,18 +522,23 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
 
         private void EnsureExecutionResourcesLocked(int requiredCapacity)
         {
+			Accelerator accelerator = _accelerator;
+
             if (_kernel == null)
             {
-                _kernel = _owner.GetKernel(_accelerator);
+                _kernel = GetKernel(_accelerator);
             }
 
             if (_divisorBuffer == null)
-            {
-                _divisorBuffer = _accelerator.Allocate1D<GpuDivisorPartialData>(1);
-                _offsetBuffer = _accelerator.Allocate1D<int>(1);
-                _countBuffer = _accelerator.Allocate1D<int>(1);
-                _cycleBuffer = _accelerator.Allocate1D<ulong>(1);
-                _firstHitBuffer = _accelerator.Allocate1D<int>(1);
+			{
+				// lock(accelerator)
+				{
+					_divisorBuffer = accelerator.Allocate1D<GpuDivisorPartialData>(1);
+					_offsetBuffer = accelerator.Allocate1D<int>(1);
+					_countBuffer = accelerator.Allocate1D<int>(1);
+					_cycleBuffer = accelerator.Allocate1D<ulong>(1);
+					_firstHitBuffer = accelerator.Allocate1D<int>(1);					
+				}
             }
 
             bool allocateBuffers = _exponentsBuffer == null;
@@ -547,10 +560,14 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
             if (allocateBuffers)
             {
                 int desiredCapacity = requiredCapacity > _capacity ? requiredCapacity : _capacity;
-                _capacity = desiredCapacity;
-                _exponentsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
-                _hitBuffer = _accelerator.Allocate1D<byte>(_capacity);
-                _hostBuffer = ThreadStaticPools.UlongPool.Rent(_capacity);
+				_capacity = desiredCapacity;
+				// lock (accelerator)
+				{
+					_exponentsBuffer = accelerator.Allocate1D<ulong>(desiredCapacity);
+					_hitBuffer = accelerator.Allocate1D<byte>(desiredCapacity);
+				}
+				
+                _hostBuffer = ThreadStaticPools.UlongPool.Rent(desiredCapacity);
             }
         }
 
@@ -769,15 +786,19 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
         }
 
         private void AllocateDeviceBuffers(Accelerator accelerator, int capacity)
-        {
-            _divisorDataBuffer = accelerator.Allocate1D<GpuDivisorPartialData>(capacity);
-            _offsetBuffer = accelerator.Allocate1D<int>(capacity);
-            _countBuffer = accelerator.Allocate1D<int>(capacity);
-            _cycleBuffer = accelerator.Allocate1D<ulong>(capacity);
-            _exponentBuffer = accelerator.Allocate1D<ulong>(capacity);
-            _hitsBuffer = accelerator.Allocate1D<byte>(capacity);
-            _hitIndexBuffer = accelerator.Allocate1D<int>(1);
-            _deviceCapacity = capacity;
+		{
+			// lock (accelerator)
+			{
+				_divisorDataBuffer = accelerator.Allocate1D<GpuDivisorPartialData>(capacity);
+				_offsetBuffer = accelerator.Allocate1D<int>(capacity);
+				_countBuffer = accelerator.Allocate1D<int>(capacity);
+				_cycleBuffer = accelerator.Allocate1D<ulong>(capacity);
+				_exponentBuffer = accelerator.Allocate1D<ulong>(capacity);
+				_hitsBuffer = accelerator.Allocate1D<byte>(capacity);
+				_hitIndexBuffer = accelerator.Allocate1D<int>(1);
+			}
+			
+			_deviceCapacity = capacity;				
         }
 
         private void RentHostArrays(int capacity)

PerfectNumbers.Core/Gpu/MersenneNumberDivisorGpuTester.cs

@@ -10,16 +10,29 @@ namespace PerfectNumbers.Core.Gpu;
 /// </summary>
 public sealed class MersenneNumberDivisorGpuTester
 {
-	private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, ReadOnlyGpuUInt128, ArrayView<byte>>> _kernelCache = new();
-	private readonly ConcurrentDictionary<Accelerator, MemoryBuffer1D<byte, Stride1D.Dense>> _resultBuffers = new();
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Queue<MemoryBuffer1D<byte, Stride1D.Dense>>>? _resultBuffers;
 
-	private Action<AcceleratorStream, Index1D, ulong, ReadOnlyGpuUInt128, ArrayView<byte>> GetKernel(Accelerator accelerator) =>
-			_kernelCache.GetOrAdd(accelerator, acc =>
-			{
-				var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ReadOnlyGpuUInt128, ArrayView<byte>>(DivisorKernels.Kernel);
-				var kernel = KernelUtil.GetKernel(loaded);
-				return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ReadOnlyGpuUInt128, ArrayView<byte>>>();
-			});
+	private static Queue<MemoryBuffer1D<byte, Stride1D.Dense>> GetQueue(Accelerator accelerator)
+	{
+		var pool = _resultBuffers ??= [];
+		if (!pool.TryGetValue(accelerator, out var bufferQueue))
+		{
+			bufferQueue = [];
+			pool[accelerator] = bufferQueue;
+		}
+
+		return bufferQueue;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ulong, ReadOnlyGpuUInt128, ArrayView<byte>> GetDivisorKernel(Accelerator accelerator)
+	{
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ReadOnlyGpuUInt128, ArrayView<byte>>(DivisorKernels.Kernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ReadOnlyGpuUInt128, ArrayView<byte>>>();
+	}
 
 	public static void BuildDivisorCandidates()
 	{
@@ -43,18 +56,30 @@ public sealed class MersenneNumberDivisorGpuTester
 
 	public bool IsDivisible(ulong exponent, in ReadOnlyGpuUInt128 divisor)
 	{
-		var accelerator = SharedGpuContext.Accelerator;
+		// var accelerator = SharedGpuContext.CreateAccelerator();
+		var accelerator = AcceleratorPool.Shared.Rent();
 		var stream = accelerator.CreateStream();
-		var kernel = GetKernel(accelerator);
-		var resultBuffer = _resultBuffers.GetOrAdd(accelerator, acc => acc.Allocate1D<byte>(1));
+		Queue<MemoryBuffer1D<byte, Stride1D.Dense>> resultBuffers = GetQueue(accelerator);
+		if (!resultBuffers.TryDequeue(out var resultBuffer))
+		{
+			// lock(accelerator)
+			{
+				resultBuffer = accelerator.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(stream);
-		kernel(stream, 1, exponent, divisor, resultBuffer.View);
-		Span<byte> result = stackalloc byte[1];
-		resultBuffer.View.CopyToCPU(stream, ref result[0], 1);
+		var divisorKernel = GetDivisorKernel(accelerator);
+		divisorKernel(stream, 1, exponent, divisor, resultBuffer.View);
+		byte result = 0;
+		resultBuffer.View.CopyToCPU(stream, ref result, 1);
 		stream.Synchronize();
-		bool divisible = result[0] != 0;
+		bool divisible = result != 0;
+		resultBuffers.Enqueue(resultBuffer);
 		stream.Dispose();
+		// AcceleratorPool.Shared.Return(accelerator);
+		// accelerator.Dispose();
 		return divisible;
 	}
 

PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs

@@ -17,14 +17,14 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
         throw new NotImplementedException($"GPU incremental scanning requires the device cycle heuristics implementation (kernel: {_kernelType}, GPU order: {_useGpuOrder}).");
 
         GpuPrimeWorkLimiter.Acquire();
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
         int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy
         UInt128 kStart = 1UL;
         ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL;
         byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; // ILGPU kernels do not support bool parameters
 
-		KernelContainer kernels = GpuKernelPool.Kernels;
+		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream);
 		var pow2Kernel = kernels.Pow2Mod!;
         var incKernel = kernels.Incremental!;
         ulong step10 = (exponent.Mod10() << 1).Mod10();
@@ -32,11 +32,11 @@ public class MersenneNumberIncrementalGpuTester(GpuKernelType kernelType, bool u
         ulong step3 = ((exponent % 3UL) << 1) % 3UL;
         ulong step5 = ((exponent % 5UL) << 1) % 5UL;
         GpuUInt128 twoPGpu = (GpuUInt128)twoP;
-        var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice();
+        var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(kernels);
         ResiduePrimeViews primeViews = default;
         if (_kernelType == GpuKernelType.Pow2Mod)
         {
-            primeViews = GpuKernelPool.GetSmallPrimesOnDevice();
+            primeViews = GpuKernelPool.GetSmallPrimesOnDevice(kernels);
         }
 
         var orderBuffer = accelerator.Allocate1D<ulong>(batchSize);

PerfectNumbers.Core/Gpu/MersenneNumberLucasLehmerGpuTester.cs

@@ -1,7 +1,6 @@
 using System.Buffers;
 using System.Collections.Concurrent;
 using ILGPU;
-using ILGPU.Algorithms;
 using ILGPU.Runtime;
 using UInt128 = System.UInt128;
 
@@ -9,60 +8,125 @@ namespace PerfectNumbers.Core.Gpu;
 
 public class MersenneNumberLucasLehmerGpuTester
 {
-    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>>> KernelCache = new();
-    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>> AddSmallKernelCache = new();
-    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>> SubSmallKernelCache = new();
-    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>> ReduceKernelCache = new();
-    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<byte>>> IsZeroKernelCache = new();
-    private readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>>> BatchKernelCache = new();
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>>>? _lucasLehmerKernel;
 
-    private Action<AcceleratorStream, Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>> GetKernel(Accelerator accelerator) =>
-        KernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>>(LucasLehmerKernels.Kernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>>>();
-        });
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>? _addSmallKernel;
 
-    private Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>> GetAddSmallKernel(Accelerator accelerator) =>
-        AddSmallKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<GpuUInt128>>(LucasLehmerKernels.AddSmallKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>();
-        });
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>? _subSmallKernel;
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>? _reduceKernel;
 
-    private Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>> GetSubSmallKernel(Accelerator accelerator) =>
-        SubSmallKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<GpuUInt128>>(LucasLehmerKernels.SubtractSmallKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>();
-        });
+	[ThreadStatic]
+    private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<byte>>>? _isZeroKernel;
 
-    private Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>> GetReduceKernel(Accelerator accelerator) =>
-        ReduceKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<GpuUInt128>>(LucasLehmerKernels.ReduceModMersenneKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>();
-        });
+	[ThreadStatic]
+    private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>>>? _lucasLehmerBatchKernel;
 
-    private Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<byte>> GetIsZeroKernel(Accelerator accelerator) =>
-        IsZeroKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<byte>>(LucasLehmerKernels.IsZeroKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<byte>>>();
-        });
+    private static Action<AcceleratorStream, Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>> GetLucasLehmerKernel(Accelerator accelerator)
+	{
+		var pool = _lucasLehmerKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
 
-    private Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>> GetBatchKernel(Accelerator accelerator) =>
-        BatchKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>>(LucasLehmerKernels.KernelBatch);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>>>();
-        });
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>>(LucasLehmerKernels.Kernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, GpuUInt128, ArrayView<GpuUInt128>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>> GetAddSmallKernel(Accelerator accelerator)
+	{
+		var pool = _addSmallKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<GpuUInt128>>(LucasLehmerKernels.AddSmallKernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>> GetSubSmallKernel(Accelerator accelerator)
+	{
+		var pool = _subSmallKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<GpuUInt128>>(LucasLehmerKernels.SubtractSmallKernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>> GetReduceKernel(Accelerator accelerator)
+	{
+		var pool = _reduceKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<GpuUInt128>>(LucasLehmerKernels.ReduceModMersenneKernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView<GpuUInt128>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<byte>> GetIsZeroKernel(Accelerator accelerator)
+	{
+		var pool = _isZeroKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<byte>>(LucasLehmerKernels.IsZeroKernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<byte>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>> GetLucasLehmerBatchKernel(Accelerator accelerator)
+	{
+		var pool = _lucasLehmerBatchKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+		
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>>(LucasLehmerKernels.KernelBatch);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
 
     // Configurable LL slice size to keep kernels short. Default 32.
     public int SliceSize = 32;
@@ -97,12 +161,12 @@ public class MersenneNumberLucasLehmerGpuTester
         }
         else
         {
-            var accelerator = SharedGpuContext.Accelerator;
+            var accelerator = AcceleratorPool.Shared.Rent();
             var stream = accelerator.CreateStream();
-            var kernel = GetKernel(accelerator);
             var modulus = new GpuUInt128(((UInt128)1 << (int)exponent) - 1UL); // TODO: Cache these Mersenne moduli per exponent so LL GPU runs skip rebuilding them every launch.
-            var buffer = accelerator.Allocate1D<GpuUInt128>(1);
-            kernel(stream, 1, exponent, modulus, buffer.View);
+			var buffer = accelerator.Allocate1D<GpuUInt128>(1);
+			var lucasLehmerKernel = GetLucasLehmerKernel(accelerator);
+            lucasLehmerKernel(stream, 1, exponent, modulus, buffer.View);
 			stream.Synchronize();
 			// TODO: Replace this with .CopyToCPU with pre-allocated buffer using stream. Synchronize on the stream after copy.
             result = buffer.GetAsArray1D()[0].IsZero;
@@ -130,9 +194,8 @@ public class MersenneNumberLucasLehmerGpuTester
             }
         }
 
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
-        var kernel = GetBatchKernel(accelerator);
 
         var expBuffer = accelerator.Allocate1D<ulong>(count);
 		ArrayPool<ulong> ulongPool = ThreadStaticPools.UlongPool;
@@ -152,8 +215,9 @@ public class MersenneNumberLucasLehmerGpuTester
         var modBuffer = accelerator.Allocate1D<GpuUInt128>(count);
         modBuffer.View.CopyFromCPU(stream, ref modulusArray[0], count);
 
-        var stateBuffer = accelerator.Allocate1D<GpuUInt128>(count);
-        kernel(stream, count, expBuffer.View, modBuffer.View, stateBuffer.View);
+		var stateBuffer = accelerator.Allocate1D<GpuUInt128>(count);
+		var lucasLehmerBatchKernel = GetLucasLehmerBatchKernel(accelerator);
+        lucasLehmerBatchKernel(stream, count, expBuffer.View, modBuffer.View, stateBuffer.View);
         stateBuffer.View.CopyToCPU(stream, ref residues[0], count);
         stream.Synchronize();
 
@@ -185,13 +249,12 @@ public class MersenneNumberLucasLehmerGpuTester
 
     private bool IsMersennePrimeNtt(ulong exponent, GpuUInt128 nttMod, GpuUInt128 primitiveRoot, bool runOnGpu)
     {
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
 
-		var addKernel = GetAddSmallKernel(accelerator);
         var subKernel = GetSubSmallKernel(accelerator);
-        var reduceKernel = GetReduceKernel(accelerator);
-        var zeroKernel = GetIsZeroKernel(accelerator);
+		var reduceKernel = GetReduceKernel(accelerator);
+		var isZeroKernel = GetIsZeroKernel(accelerator);
 
         int limbCount = (int)((exponent + 127UL) / 128UL);
         int target = limbCount * 2;
@@ -201,8 +264,9 @@ public class MersenneNumberLucasLehmerGpuTester
             transformLength <<= 1;
         }
         var stateBuffer = accelerator.Allocate1D<GpuUInt128>(transformLength);
-        stateBuffer.MemSetToZero(stream);
-        addKernel(stream, 1, 4UL, stateBuffer.View);
+		stateBuffer.MemSetToZero(stream);
+		var addSmallKernel = GetAddSmallKernel(accelerator);
+		addSmallKernel(stream, 1, 4UL, stateBuffer.View);
 
         // TODO(LL-SLICE): Slice Lucas–Lehmer iterations into short batches to
         // avoid long-running kernels and TDR. Example: process 8–64 iterations
@@ -227,11 +291,12 @@ public class MersenneNumberLucasLehmerGpuTester
         }
 
         reduceKernel(stream, 1, exponent, stateBuffer.View);
-        var resultBuffer = accelerator.Allocate1D<byte>(1);
-        zeroKernel(stream, 1, stateBuffer.View, resultBuffer.View);
-        byte[] result = new byte[1];
+		var resultBuffer = accelerator.Allocate1D<byte>(1);
+        isZeroKernel(stream, 1, stateBuffer.View, resultBuffer.View);
+        Span<byte> result = stackalloc byte[1];
         resultBuffer.View.CopyToCPU(stream, result);
-        stream.Synchronize();
+		stream.Synchronize();
+		
         bool isPrime = result[0] != 0;
         resultBuffer.Dispose();
         stateBuffer.Dispose();
@@ -504,6 +569,4 @@ public class MersenneNumberLucasLehmerGpuTester
 
         return result;
     }
-
-
 }

PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs

@@ -10,28 +10,34 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
 	public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime)
 	{
 		GpuPrimeWorkLimiter.Acquire();
-		var accelerator = SharedGpuContext.Accelerator;
+		// var accelerator = SharedGpuContext.Accelerator;
+		var accelerator = AcceleratorPool.Shared.Rent();
 		var stream = accelerator.CreateStream();
 		int batchSize = GpuConstants.ScanBatchSize; // large batch improves GPU occupancy and avoids TDR
 		UInt128 kStart = UInt128.One;
 		ulong divMul = (ulong)((((UInt128)1 << 64) - UInt128.One) / exponent) + 1UL;
 		byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0; // ILGPU kernels do not support bool parameters
 
-		KernelContainer kernels = GpuKernelPool.Kernels;
+		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream);
 		var kernel = _kernelType switch
 		{
 			GpuKernelType.Pow2Mod => kernels.Pow2ModOrder!,
 			_ => kernels.IncrementalOrder!,
 		};
 
-		var foundBuffer = accelerator.Allocate1D<int>(1); // TODO: Replace this allocation with the pooled device buffer from GpuOrderKernelBenchmarks so repeated scans reuse the pinned staging memory instead of allocating per run.
+		MemoryBuffer1D<int, ILGPU.Stride1D.Dense>? foundBuffer;
+		// lock (accelerator)
+		{
+			foundBuffer = accelerator.Allocate1D<int>(1); // TODO: Replace this allocation with the pooled device buffer from GpuOrderKernelBenchmarks so repeated scans reuse the pinned staging memory instead of allocating per run.
+
+		}
 		exponent.Mod10_8_5_3Steps(out ulong step10, out ulong step8, out ulong step5, out ulong step3);
 
 		UInt128 remaining;
 		int currentSize;
 		int found = 0;
 
-		var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice();
+		var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(kernels);
 		while (kStart <= maxK && Volatile.Read(ref isPrime))
 		{
 			remaining = maxK - kStart + 1UL;

PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs

@@ -1,10 +1,5 @@
-using PerfectNumbers.Core.Gpu;
-using System;
-using System.Buffers;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Runtime.InteropServices;
-using System.Threading;
 using System.Runtime.CompilerServices;
 using ILGPU;
 using ILGPU.Runtime;
@@ -13,174 +8,208 @@ namespace PerfectNumbers.Core.Gpu;
 
 public class MersenneNumberResidueGpuTester(bool useGpuOrder)
 {
-    private readonly bool _useGpuOrder = useGpuOrder;
-    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);
+	private readonly bool _useGpuOrder = useGpuOrder;
 
-    // TODO: Integrate MersenneDivisorCycles.Shared to consult cycle lengths for small q (<= 4M) and fast-reject
-    // candidates without launching heavy kernels. Consider a device-visible constant buffer per-accelerator
-    // via GpuKernelPool to avoid host round-trips.
-    // TODO: When cycles are missing for larger q values, compute only the single required cycle on the device
-    // requested by the caller, skip queuing additional block generation, and keep the snapshot cache untouched.
+	[ThreadStatic]
+	private static Dictionary<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>>>? _pow2ModKernel;
 
-    // GPU residue variant: check 2^p % q == 1 for q = 2*p*k + 1.
-    public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime)
-    {
-        GpuPrimeWorkLimiter.Acquire();
-        var accelerator = SharedGpuContext.Accelerator;
+	[ThreadStatic]
+	private static Queue<ResidueResources>? _resourcePools;
 
-        // Monitor.Enter(gpuLease.ExecutionLock);
+	// TODO: Integrate MersenneDivisorCycles.Shared to consult cycle lengths for small q (<= 4M) and fast-reject
+	// candidates without launching heavy kernels. Consider a device-visible constant buffer per-accelerator
+	// via GpuKernelPool to avoid host round-trips.
+	// TODO: When cycles are missing for larger q values, compute only the single required cycle on the device
+	// requested by the caller, skip queuing additional block generation, and keep the snapshot cache untouched.
 
-        var stream = accelerator.CreateStream();
-        var kernel = GetPow2ModKernel(accelerator);
+	// GPU residue variant: check 2^p % q == 1 for q = 2*p*k + 1.
+	public void Scan(ulong exponent, UInt128 twoP, LastDigit lastDigit, UInt128 maxK, ref bool isPrime)
+	{
+		GpuPrimeWorkLimiter.Acquire();
+		var accelerator = AcceleratorPool.Shared.Rent();
+
+		// Monitor.Enter(gpuLease.ExecutionLock);
+
+		var stream = accelerator.CreateStream();
+		var kernel = GetPow2ModKernel(accelerator);
 		var resources = RentResources(accelerator, GpuConstants.ScanBatchSize);
-		var gpuKernels = GpuKernelPool.Kernels;
+		var gpuKernels = GpuKernelPool.GetOrAddKernels(accelerator, stream);
 
-        var orderBuffer = resources.OrderBuffer;
-        ulong[] orderArray = resources.OrderArray;
+		var orderBuffer = resources.OrderBuffer;
+		ulong[] orderArray = resources.OrderArray;
 		int batchSize = resources.Capacity;
 
-        // Ensure device has small cycles and primes tables for in-kernel lookup
+		// Ensure device has small cycles and primes tables for in-kernel lookup
 		// TODO: Pass-in the stream down the path
-        var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice();
-        ResiduePrimeViews primeViews = GpuKernelPool.GetSmallPrimesOnDevice();
+		var smallCyclesView = GpuKernelPool.GetSmallCyclesOnDevice(gpuKernels);
+		ResiduePrimeViews primeViews = GpuKernelPool.GetSmallPrimesOnDevice(gpuKernels);
 
-        UInt128 kStart = 1UL;
-        UInt128 limit = maxK + UInt128.One;
-        byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0;
-        twoP.Mod10_8_5_3(out ulong step10, out ulong step8, out ulong step5, out ulong step3);
-        step10 = step10.Mod10();
+		UInt128 kStart = 1UL;
+		UInt128 limit = maxK + UInt128.One;
+		byte last = lastDigit == LastDigit.Seven ? (byte)1 : (byte)0;
+		twoP.Mod10_8_5_3(out ulong step10, out ulong step8, out ulong step5, out ulong step3);
+		step10 = step10.Mod10();
 
-        GpuUInt128 twoPGpu = (GpuUInt128)twoP;
-        UInt128 batchSize128 = (UInt128)batchSize;
-        UInt128 fullBatchStep = twoP * batchSize128;
-        UInt128 q = twoP * kStart + UInt128.One;
+		GpuUInt128 twoPGpu = (GpuUInt128)twoP;
+		UInt128 batchSize128 = (UInt128)batchSize;
+		UInt128 fullBatchStep = twoP * batchSize128;
+		UInt128 q = twoP * kStart + UInt128.One;
 
-        while (kStart < limit && Volatile.Read(ref isPrime))
-        {
-            UInt128 remaining = limit - kStart;
-            int currentSize = remaining > batchSize128 ? batchSize : (int)remaining;
-            Span<ulong> orders = orderArray.AsSpan(0, currentSize);
-            ref ulong ordersRef = ref MemoryMarshal.GetReference(orders);
+		while (kStart < limit && Volatile.Read(ref isPrime))
+		{
+			UInt128 remaining = limit - kStart;
+			int currentSize = remaining > batchSize128 ? batchSize : (int)remaining;
+			Span<ulong> orders = orderArray.AsSpan(0, currentSize);
+			ref ulong ordersRef = ref MemoryMarshal.GetReference(orders);
 
-            q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
-            var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
+			q.Mod10_8_5_3(out ulong q0m10, out ulong q0m8, out ulong q0m5, out ulong q0m3);
+			var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
 
-            kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
-                kernelArgs, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2);
+			kernel(stream, currentSize, exponent, twoPGpu, (GpuUInt128)kStart, last, 0UL,
+				kernelArgs, orderBuffer.View, smallCyclesView, primeViews.LastOne, primeViews.LastSeven, primeViews.LastOnePow2, primeViews.LastSevenPow2);
 
-			
-            orderBuffer.View.CopyToCPU(stream, ref ordersRef, currentSize);
+			orderBuffer.View.CopyToCPU(stream, ref ordersRef, currentSize);
 			stream.Synchronize();
-            if (!Volatile.Read(ref isPrime))
-            {
-                break;
-            }
+			if (!Volatile.Read(ref isPrime))
+			{
+				break;
+			}
 
-            bool compositeFound = false;
-            for (int i = 0; i < currentSize; i++)
-            {
-                if (orders[i] != 0UL)
-                {
-                    Volatile.Write(ref isPrime, false);
-                    compositeFound = true;
-                    break;
-                }
-            }
+			bool compositeFound = false;
+			for (int i = 0; i < currentSize; i++)
+			{
+				if (orders[i] != 0UL)
+				{
+					Volatile.Write(ref isPrime, false);
+					compositeFound = true;
+					break;
+				}
+			}
 
-            if (compositeFound)
-            {
-                break;
-            }
+			if (compositeFound)
+			{
+				break;
+			}
 
-            UInt128 processed = (UInt128)currentSize;
-            kStart += processed;
-            if (kStart < limit)
-            {
-                if (processed == batchSize128)
-                {
-                    q += fullBatchStep;
-                }
-                else
-                {
-                    q += twoP * processed; // TODO: Swap this multiplication for the shared cycle stepping helper once residue cycles are mandatory so we can reuse the cached remainder ladder instead of performing UInt128 multiply operations.
-                }
-            }
-        }
+			UInt128 processed = (UInt128)currentSize;
+			kStart += processed;
+			if (kStart < limit)
+			{
+				if (processed == batchSize128)
+				{
+					q += fullBatchStep;
+				}
+				else
+				{
+					q += twoP * processed; // TODO: Swap this multiplication for the shared cycle stepping helper once residue cycles are mandatory so we can reuse the cached remainder ladder instead of performing UInt128 multiply operations.
+				}
+			}
+		}
 
-        ReturnResources(accelerator, resources);
-        stream.Dispose();
+		ReturnResources(resources);
+		stream.Dispose();
 		// Monitor.Exit(gpuLease.ExecutionLock);
 		gpuKernels.Dispose();
 		GpuPrimeWorkLimiter.Release();
-    }
+	}
 
-    private 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>> GetPow2ModKernel(Accelerator accelerator) =>
-        _pow2ModKernelCache.GetOrAdd(
-            accelerator,
-            static acc =>
-            {
-                var loaded = acc.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);
-                return (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>>)Delegate.CreateDelegate(
-                    typeof(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>>),
-                    loaded.Target,
-                    loaded.Method);
-            });
+	private static 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>> GetPow2ModKernel(Accelerator accelerator)
+	{
 
-    private ResidueResources RentResources(Accelerator accelerator, int capacity)
-    {
-        var bag = _resourcePools.GetOrAdd(accelerator, static _ => []);
-        while (bag.TryTake(out var resources))
-        {
-            if (resources.Capacity >= capacity)
+		var pool = _pow2ModKernel ?? [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		cached = CreatePow2ModKernel(accelerator);
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static 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>> CreatePow2ModKernel(Accelerator accelerator)
+	{
+		var loaded = accelerator.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 = (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>>)Delegate.CreateDelegate(
+			typeof(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>>),
+			loaded.Target,
+			loaded.Method);
+		return kernel;
+	}
+
+
+	private static ResidueResources RentResources(Accelerator accelerator, int capacity)
+	{
+		var resourcePools = _resourcePools ??= new();
+		if (resourcePools.TryDequeue(out var bag))
+		{
+			bag.EnsureCapacity(accelerator, capacity);
+			return bag;
+		}
+
+		return new ResidueResources(accelerator, capacity);
+	}
+
+	private static void ReturnResources(ResidueResources resources)
+	{
+		var resourcePools = _resourcePools ??= new();
+		resourcePools.Enqueue(resources);
+	}
+
+	private sealed class ResidueResources
+	{
+		internal ResidueResources(Accelerator accelerator, int capacity)
+		{
+			Capacity = Math.Max(1, capacity);
+			// lock (accelerator)
 			{
-				// The caller always overwrites the required elements. We don't need to worry about clearing the output buffer.
-                // resources.OrderBuffer.MemSetToZero();
-                return resources;
-            }
+				OrderBuffer = accelerator.Allocate1D<ulong>(Capacity);
+			}
 
-            resources.Dispose();
-        }
+			OrderArray = ThreadStaticPools.UlongPool.Rent(Capacity);
+		}
 
-        var created = new ResidueResources(accelerator, capacity);
-		// The caller always overwrites the required elements. We don't need to worry about clearing the output buffer.
-        // created.OrderBuffer.MemSetToZero();
-        return created;
-    }
+		internal void EnsureCapacity(Accelerator accelerator, int capacity)
+		{
+			if (OrderBuffer.Length >= capacity)
+			{
+				return;
+			}
 
-    private void ReturnResources(Accelerator accelerator, ResidueResources resources) =>
-        _resourcePools.GetOrAdd(accelerator, static _ => new ConcurrentBag<ResidueResources>()).Add(resources);
+			OrderBuffer.Dispose();
+			ThreadStaticPools.UlongPool.Return(OrderArray);
 
-    private sealed class ResidueResources
-    {
-        internal ResidueResources(Accelerator accelerator, int capacity)
-        {
-            Capacity = Math.Max(1, capacity);
-            OrderBuffer = accelerator.Allocate1D<ulong>(Capacity);
-            OrderArray = ThreadStaticPools.UlongPool.Rent(Capacity);
-        }
+			// lock (accelerator)
+			{
+				OrderBuffer = accelerator.Allocate1D<ulong>(capacity);
+			}
 
-        internal MemoryBuffer1D<ulong, Stride1D.Dense> OrderBuffer { get; }
+			OrderArray = ThreadStaticPools.UlongPool.Rent(capacity);
+			Capacity = capacity;
+		}
 
-        internal ulong[] OrderArray { get; }
+		internal MemoryBuffer1D<ulong, Stride1D.Dense> OrderBuffer { get; private set; }
 
-        internal int Capacity { get; }
+		internal ulong[] OrderArray { get; private set; }
 
-        public void Dispose()
-        {
-            OrderBuffer.Dispose();
-            ThreadStaticPools.UlongPool.Return(OrderArray, clearArray: false);
-        }
-    }
+		internal int Capacity { get; private set; }
 
-    private sealed class AcceleratorReferenceComparer : IEqualityComparer<Accelerator>
-    {
-        public static AcceleratorReferenceComparer Instance { get; } = new();
+		public void Dispose()
+		{
+			OrderBuffer.Dispose();
+			ThreadStaticPools.UlongPool.Return(OrderArray, clearArray: false);
+		}
+	}
 
-        public bool Equals(Accelerator? x, Accelerator? y) => ReferenceEquals(x, y);
+	private sealed class AcceleratorReferenceComparer : IEqualityComparer<Accelerator>
+	{
+		public static AcceleratorReferenceComparer Instance { get; } = new();
 
-        public int GetHashCode(Accelerator obj) => RuntimeHelpers.GetHashCode(obj);
-    }
+		public bool Equals(Accelerator? x, Accelerator? y) => ReferenceEquals(x, y);
+
+		public int GetHashCode(Accelerator obj) => RuntimeHelpers.GetHashCode(obj);
+	}
 
 
 }

PerfectNumbers.Core/Gpu/NttGpuMath.cs

@@ -29,41 +29,62 @@ public static class NttGpuMath
     public static NttBackend GpuTransformBackend { get; set; } = NttBackend.Reference;
 	// Controls modular reduction strategy inside staged NTT paths.
 	public static ModReductionMode ReductionMode { get; set; } = ModReductionMode.Auto;
-	
-    private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128>> MulKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream,Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>> StageKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
 
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128>> ScaleKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128>>? _mulKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
 	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong>> StageMontKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong>> StageBarrett128KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>> ToMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong>> FromMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong>> ScaleBarrett128KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>> SquareMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>> ScaleMont64KernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128>> ForwardKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128>> InverseKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-	private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int>> BitReverseKernelCache = new(); // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
-	
-    private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128> GetStageKernel(Accelerator accelerator) =>
-        StageKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>(NttButterflyKernels.StageKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>>();
-        });
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream,Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>>? _stageKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
 
-    private sealed class SquareCacheEntry
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128>>? _scaleKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong>>? _stageMontKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong>>? _stageBarrett128Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>>? _toMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong>>? _fromMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong>>? _scaleBarrett128Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static readonly Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>>? _squareMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>>? _scaleMont64Kernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128>>? _forwardKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128>>? _inverseKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+	
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int>>? _bitReverseKernel; // TODO: Replace this concurrent cache with the prewarmed accelerator-indexed tables from GpuModularArithmeticBenchmarks so kernel launches avoid dictionary lookups once the kernels are baked during startup.
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128> GetStageKernel(Accelerator accelerator)
+	{
+		var pool = _stageKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+		
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>(NttButterflyKernels.StageKernel);
+		var kernel = KernelUtil.GetKernel(loaded);
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, GpuUInt128>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private sealed class SquareCacheEntry
     {
         public MemoryBuffer1D<GpuUInt128, Stride1D.Dense> Buffer { get; }
         public GpuUInt128 Root { get; }
@@ -266,7 +287,7 @@ public static class NttGpuMath
 
     private readonly record struct SquareCacheKey(int Length, GpuUInt128 Modulus, GpuUInt128 PrimitiveRoot);
 
-    private static readonly ConcurrentDictionary<Accelerator, ConcurrentDictionary<SquareCacheKey, SquareCacheEntry>> SquareCache = new();
+    private static readonly ConcurrentDictionary<SquareCacheKey, SquareCacheEntry> SquareCache = [];
 
     // TODO(NTT-OPT): Add a TwiddleCache for stage-wise Cooley–Tukey NTT.
     // Design:
@@ -276,9 +297,8 @@ public static class NttGpuMath
 
     private static SquareCacheEntry GetSquareCache(Accelerator accelerator, AcceleratorStream stream, int length, GpuUInt128 modulus, GpuUInt128 primitiveRoot)
     {
-        var perAccel = SquareCache.GetOrAdd(accelerator, _ => new ConcurrentDictionary<SquareCacheKey, SquareCacheEntry>());
         var key = new SquareCacheKey(length, modulus, primitiveRoot);
-        return perAccel.GetOrAdd(key, _ => new SquareCacheEntry(accelerator, stream, length, modulus, primitiveRoot));
+        return SquareCache.GetOrAdd(key, _ => new SquareCacheEntry(accelerator, stream, length, modulus, primitiveRoot));
     }
 
     private static readonly ulong[] SmallPrimes = GenerateSmallPrimes(1000);
@@ -346,147 +366,197 @@ public static class NttGpuMath
     {
         if (accelerator is null)
         {
-            foreach (var perAccel in SquareCache.Values)
+            foreach (var entry in SquareCache.Values)
             {
-                foreach (var entry in perAccel.Values)
-                {
-                    entry.Dispose();
-                }
+				entry.Dispose();
             }
 
             SquareCache.Clear();
-            MulKernelCache.Clear();
-            StageKernelCache.Clear();
-            ScaleKernelCache.Clear();
-            StageMontKernelCache.Clear();
-            StageBarrett128KernelCache.Clear();
-            ToMont64KernelCache.Clear();
-            FromMont64KernelCache.Clear();
-            SquareMont64KernelCache.Clear();
-            ScaleMont64KernelCache.Clear();
-            ForwardKernelCache.Clear();
-			InverseKernelCache.Clear();
-			BitReverseKernelCache.Clear();
             // TODO(NTT-OPT): Clear TwiddleCache once introduced.
             return;
         }
 
-        if (SquareCache.TryRemove(accelerator, out var cache))
-        {
-            foreach (var entry in cache.Values)
-            {
-                entry.Dispose();
-            }
-        }
+		var squareCache = SquareCache; 
+		foreach (var entry in squareCache.Values)
+		{
+			entry.Dispose();
+		}
 
-        MulKernelCache.TryRemove(accelerator, out _);
-        StageKernelCache.TryRemove(accelerator, out _);
-        ScaleKernelCache.TryRemove(accelerator, out _);
-        StageMontKernelCache.TryRemove(accelerator, out _);
-        StageBarrett128KernelCache.TryRemove(accelerator, out _);
-        ToMont64KernelCache.TryRemove(accelerator, out _);
-        FromMont64KernelCache.TryRemove(accelerator, out _);
-        SquareMont64KernelCache.TryRemove(accelerator, out _);
-        ScaleMont64KernelCache.TryRemove(accelerator, out _);
-        ForwardKernelCache.TryRemove(accelerator, out _);
-        InverseKernelCache.TryRemove(accelerator, out _);
-        BitReverseKernelCache.TryRemove(accelerator, out _);
         // TODO(NTT-OPT): Remove per-accelerator twiddle buffers once added.
     }
 
-    private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int> GetBitReverseKernel(Accelerator accelerator) =>
-        BitReverseKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int>(NttTransformKernels.BitReverseKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int>>();
-        });
-
-    private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128> GetForwardKernel(Accelerator accelerator) =>
-        ForwardKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128>(NttTransformKernels.ForwardKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128>>();
-        });
-
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong> GetFromMont64Kernel(Accelerator accelerator) =>
-		FromMont64KernelCache.GetOrAdd(accelerator, acc =>
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int> GetBitReverseKernel(Accelerator accelerator)
+	{
+		var pool = _bitReverseKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ulong, ulong>(NttMontgomeryKernels.FromMont64Kernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong>>();
-		});
+			return cached;
+		}
 
-    private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128> GetInverseKernel(Accelerator accelerator) =>
-        InverseKernelCache.GetOrAdd(accelerator, acc =>
-        {
-            var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128>(NttTransformKernels.InverseKernel);
-            var kernel = KernelUtil.GetKernel(loaded);
-            return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128>>();
-        });
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int>(NttTransformKernels.BitReverseKernel);
+		
+		var kernel = KernelUtil.GetKernel(loaded);
 
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128> GetMulKernel(Accelerator accelerator) =>
-		MulKernelCache.GetOrAdd(accelerator, valueFactory: acc =>
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128> GetForwardKernel(Accelerator accelerator)
+	{
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128>(NttTransformKernels.ForwardKernel);
+		var kernel = KernelUtil.GetKernel(loaded);
+		return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128>>();
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong> GetFromMont64Kernel(Accelerator accelerator)
+	{
+		var pool = _fromMont64Kernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128>(NttPointwiseKernels.MulKernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128>>();
-		});
+			return cached;
+		}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128> GetScaleKernel(Accelerator accelerator) =>
-		ScaleKernelCache.GetOrAdd(accelerator, acc =>
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ulong, ulong>(NttMontgomeryKernels.FromMont64Kernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128> GetInverseKernel(Accelerator accelerator)
+	{
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128>(NttTransformKernels.InverseKernel);
+		var kernel = KernelUtil.GetKernel(loaded);
+		return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, int, GpuUInt128, GpuUInt128, GpuUInt128>>();
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128> GetMulKernel(Accelerator accelerator)
+	{
+		var pool = _mulKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128>(NttPointwiseKernels.ScaleKernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128>>();
-		});
+			return cached;
+		}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong> GetScaleBarrett128Kernel(Accelerator accelerator) =>
-		ScaleBarrett128KernelCache.GetOrAdd(accelerator, acc =>
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128>(NttPointwiseKernels.MulKernel);
+		
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ArrayView<GpuUInt128>, GpuUInt128>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128> GetScaleKernel(Accelerator accelerator)
+	{
+		var pool = _scaleKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong>(NttBarrettKernels.ScaleBarrett128Kernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong>>();
-		});
+			return cached;
+		}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong> GetScaleMont64Kernel(Accelerator accelerator) =>
-		ScaleMont64KernelCache.GetOrAdd(accelerator, acc =>
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128>(NttPointwiseKernels.ScaleKernel);
+		
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, GpuUInt128>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong> GetScaleBarrett128Kernel(Accelerator accelerator)
+	{
+		var pool = _scaleBarrett128Kernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>(NttMontgomeryKernels.ScaleMont64Kernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>>();
-		});
+			return cached;
+		}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong> GetStageBarrett128Kernel(Accelerator accelerator) =>
-		StageBarrett128KernelCache.GetOrAdd(accelerator, acc =>
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong>(NttBarrettKernels.ScaleBarrett128Kernel);
+		
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, GpuUInt128, ulong, ulong, ulong, ulong>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong> GetScaleMont64Kernel(Accelerator accelerator)
+	{
+		var pool = _scaleMont64Kernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong>(NttBarrettKernels.StageBarrett128Kernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong>>();
-		});
+			return cached;
+		}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong> GetStageMontKernel(Accelerator accelerator) =>
-		StageMontKernelCache.GetOrAdd(accelerator, acc =>
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>(NttMontgomeryKernels.ScaleMont64Kernel);
+		
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong> GetStageBarrett128Kernel(Accelerator accelerator)
+	{
+		var pool = _stageBarrett128Kernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong>(NttMontgomeryKernels.StageMontKernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong>>();
-		});
+			return cached;
+		}
 
-	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong> GetToMont64Kernel(Accelerator accelerator) =>
-		ToMont64KernelCache.GetOrAdd(accelerator, acc =>
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong>(NttBarrettKernels.StageBarrett128Kernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+		
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong, ulong, ulong>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong> GetStageMontKernel(Accelerator accelerator)
+	{
+		var pool = _stageMontKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
 		{
-			var loaded = acc.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int>(NttTransformKernels.BitReverseKernel);
-			var kernel = KernelUtil.GetKernel(loaded);
-			return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>>();
-		});
+			return cached;
+		}
 
-    // TODO(NTT-OPT): Mirror the forward stage-wise design for the inverse
-    // transform. Precompute inverse twiddles and a normalization factor
-    // (n^-1 mod m). Launch one short kernel per stage to avoid TDR.
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong>(NttMontgomeryKernels.StageMontKernel);
 
-    public static void BitReverse(Span<GpuUInt128> values)
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, int, int, int, ArrayView<GpuUInt128>, ulong, ulong>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	private static Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong> GetToMont64Kernel(Accelerator accelerator)
+	{
+		var pool = _toMont64Kernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>(NttMontgomeryKernels.ToMont64Kernel);
+		
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<GpuUInt128>, ulong, ulong, ulong>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
+	// TODO(NTT-OPT): Mirror the forward stage-wise design for the inverse
+	// transform. Precompute inverse twiddles and a normalization factor
+	// (n^-1 mod m). Launch one short kernel per stage to avoid TDR.
+
+	public static void BitReverse(Span<GpuUInt128> values)
     {
         int n = values.Length;
         int bits = (int)Math.Log2(n);
@@ -504,15 +574,15 @@ public static class NttGpuMath
     {
         int n = values.Length;
         int bits = (int)Math.Log2(n);
-        var accelerator = SharedGpuContext.Accelerator;
-        var kernel = GetBitReverseKernel(accelerator);
+        var accelerator = AcceleratorPool.Shared.Rent();
         var pool = ThreadStaticPools.GpuUInt128Pool;
         var array = pool.Rent(n);
         values.CopyTo(array);
         var buffer = accelerator.Allocate1D<GpuUInt128>(n);
         var stream = accelerator.CreateStream();
-        buffer.View.CopyFromCPU(stream, ref array[0], n);
-        kernel(stream, n, buffer.View, bits);
+		buffer.View.CopyFromCPU(stream, ref array[0], n);
+		var bitReverseKernel = GetBitReverseKernel(accelerator);
+        bitReverseKernel(stream, n, buffer.View, bits);
         buffer.View.CopyToCPU(stream, ref array[0], n);
         stream.Synchronize();
         array.AsSpan(0, n).CopyTo(values);
@@ -566,9 +636,8 @@ public static class NttGpuMath
     private static void ForwardGpuReference(Span<GpuUInt128> values, GpuUInt128 modulus, GpuUInt128 primitiveRoot)
     {
         int n = (int)values.Length;
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
-        var kernel = GetForwardKernel(accelerator);
         var pool = ThreadStaticPools.GpuUInt128Pool;
         var inputArray = pool.Rent(n);
         var outputArray = pool.Rent(n);
@@ -579,8 +648,9 @@ public static class NttGpuMath
         UInt128 modValue = modulus;
         UInt128 expBase = (modValue - 1UL) / (ulong)n;
         var root = new GpuUInt128(primitiveRoot);
-        root.ModPow((ulong)expBase, modulus);
-        kernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, root);
+		root.ModPow((ulong)expBase, modulus);
+		var forwardKernel = GetForwardKernel(accelerator);
+        forwardKernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, root);
         outputBuffer.View.CopyToCPU(stream, ref outputArray[0], n);
         stream.Synchronize();
         outputArray.AsSpan(0, n).CopyTo(values);
@@ -662,9 +732,8 @@ public static class NttGpuMath
     private static void InverseGpuReference(Span<GpuUInt128> values, GpuUInt128 modulus, GpuUInt128 primitiveRoot)
     {
         int n = values.Length;
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
-        var kernel = GetInverseKernel(accelerator);
         var pool = ThreadStaticPools.GpuUInt128Pool;
         var inputArray = pool.Rent(n);
         var outputArray = pool.Rent(n);
@@ -687,15 +756,17 @@ public static class NttGpuMath
         }
 
         var nInv = new GpuUInt128((UInt128)n);
-        if (modulus.High == 0UL)
-        {
-            nInv.ModInv(modulus.Low);
-        }
-        else
-        {
-            nInv.ModInv(modulus);
-        }
-        kernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, rootInv, nInv);
+		if (modulus.High == 0UL)
+		{
+			nInv.ModInv(modulus.Low);
+		}
+		else
+		{
+			nInv.ModInv(modulus);
+		}
+
+		var inverseKernel = GetInverseKernel(accelerator);
+        inverseKernel(stream, n, inputBuffer.View, outputBuffer.View, n, modulus, rootInv, nInv);
         outputBuffer.View.CopyToCPU(stream, ref outputArray[0], n);
         stream.Synchronize();
         outputArray.AsSpan(0, n).CopyTo(values);
@@ -709,7 +780,7 @@ public static class NttGpuMath
     private static void ForwardGpuStaged(Span<GpuUInt128> values, GpuUInt128 modulus, GpuUInt128 primitiveRoot)
     {
         int n = values.Length;
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
         var buffer = accelerator.Allocate1D<GpuUInt128>(n);
         buffer.View.CopyFromCPU(stream, ref values[0], n);
@@ -724,7 +795,7 @@ public static class NttGpuMath
     private static void InverseGpuStaged(Span<GpuUInt128> values, GpuUInt128 modulus, GpuUInt128 primitiveRoot)
     {
         int n = values.Length;
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
         var buffer = accelerator.Allocate1D<GpuUInt128>(n);
         buffer.View.CopyFromCPU(stream, ref values[0], n);
@@ -737,24 +808,24 @@ public static class NttGpuMath
     }
 
     private static void ForwardDevice(Accelerator accelerator, AcceleratorStream stream, ArrayView<GpuUInt128> data, int length, GpuUInt128 modulus, SquareCacheEntry cache)
-    {
-        var bitKernel = GetBitReverseKernel(accelerator);
-        bitKernel(stream, length, data, cache.Bits);
+	{
+		var bitReverseKernel = GetBitReverseKernel(accelerator);
+        bitReverseKernel(stream, length, data, cache.Bits);
 
         if (cache.UseMontgomery64)
-        {
-            var toMont = GetToMont64Kernel(accelerator);
-            toMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64);
+		{
+			var toMont64Kernel = GetToMont64Kernel(accelerator);
+            toMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64);
             var stageKernel = GetStageMontKernel(accelerator);
-            int butterflies = length >> 1;
-            for (int s = 0; s < cache.Bits; s++)
-            {
-                int len = 1 << (s + 1);
-                int half = len >> 1;
-                stageKernel( stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesMont!.View, cache.ModulusLow, cache.MontNPrime64);
-            }
-            var fromMont = GetFromMont64Kernel(accelerator);
-            fromMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64);
+			int butterflies = length >> 1;
+			for (int s = 0; s < cache.Bits; s++)
+			{
+				int len = 1 << (s + 1);
+				int half = len >> 1;
+				stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesMont!.View, cache.ModulusLow, cache.MontNPrime64);
+			}
+			var fromMont64Kernel = GetFromMont64Kernel(accelerator);
+            fromMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64);
         }
         else if (cache.UseBarrett128)
         {
@@ -783,51 +854,55 @@ public static class NttGpuMath
     }
 
     private static void InverseDevice(Accelerator accelerator, AcceleratorStream stream, ArrayView<GpuUInt128> data, int length, GpuUInt128 modulus, SquareCacheEntry cache)
-    {
-        var bitKernel = GetBitReverseKernel(accelerator);
-        bitKernel(stream, length, data, cache.Bits);
+	{
+		var bitReverseKernel = GetBitReverseKernel(accelerator);
+        bitReverseKernel(stream, length, data, cache.Bits);
 
         if (cache.UseMontgomery64)
-        {
-            var toMont = GetToMont64Kernel(accelerator);
-            toMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64);
+		{
+			var toMont64Kernel = GetToMont64Kernel(accelerator);
+            toMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontR2Mod64);
             var stageKernel = GetStageMontKernel(accelerator);
             int butterflies = length >> 1;
-            for (int s = 0; s < cache.Bits; s++)
-            {
-                int len = 1 << (s + 1);
-                int half = len >> 1;
-                stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInvMont!.View, cache.ModulusLow, cache.MontNPrime64);
-            }
-            var scaleKernel = GetScaleMont64Kernel(accelerator);
-            scaleKernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontNInvR64);
-            var fromMont = GetFromMont64Kernel(accelerator);
-            fromMont(stream, length, data, cache.ModulusLow, cache.MontNPrime64);
+			for (int s = 0; s < cache.Bits; s++)
+			{
+				int len = 1 << (s + 1);
+				int half = len >> 1;
+				stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInvMont!.View, cache.ModulusLow, cache.MontNPrime64);
+			}
+
+			var scaleMont64Kernel = GetScaleMont64Kernel(accelerator);
+			scaleMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64, cache.MontNInvR64);
+
+			var fromMont64Kernel = GetFromMont64Kernel(accelerator);
+            fromMont64Kernel(stream, length, data, cache.ModulusLow, cache.MontNPrime64);
         }
         else if (cache.UseBarrett128)
         {
             var stageKernel = GetStageBarrett128Kernel(accelerator);
             int butterflies = length >> 1;
-            for (int s = 0; s < cache.Bits; s++)
-            {
-                int len = 1 << (s + 1);
-                int half = len >> 1;
-                stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow);
-            }
-            var scaleKernel = GetScaleBarrett128Kernel(accelerator);
-            scaleKernel(stream, length, data, cache.NInv, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow);
+			for (int s = 0; s < cache.Bits; s++)
+			{
+				int len = 1 << (s + 1);
+				int half = len >> 1;
+				stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow);
+			}
+
+			var scaleBarrett128Kernel = GetScaleBarrett128Kernel(accelerator);
+            scaleBarrett128Kernel(stream, length, data, cache.NInv, cache.ModulusHigh, cache.ModulusLow, cache.BarrettMuHigh, cache.BarrettMuLow);
         }
         else
         {
             var stageKernel = GetStageKernel(accelerator);
             int butterflies = length >> 1;
-            for (int s = 0; s < cache.Bits; s++)
-            {
-                int len = 1 << (s + 1);
-                int half = len >> 1;
-                stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, modulus);
-            }
-            var scaleKernel = GetScaleKernel(accelerator);
+			for (int s = 0; s < cache.Bits; s++)
+			{
+				int len = 1 << (s + 1);
+				int half = len >> 1;
+				stageKernel(stream, butterflies, data, len, half, cache.StageOffsets[s], cache.TwiddlesInv.View, modulus);
+			}
+
+			var scaleKernel = GetScaleKernel(accelerator);
             scaleKernel(stream, length, data, cache.NInv, modulus);
         }
 
@@ -917,17 +992,16 @@ public static class NttGpuMath
     public static void SquareDevice(Accelerator accelerator, AcceleratorStream stream, ArrayView<GpuUInt128> values, GpuUInt128 modulus, GpuUInt128 primitiveRoot)
     {
         var cache = GetSquareCache(accelerator, stream, (int)values.Length, modulus, primitiveRoot);
-        ForwardDevice(accelerator, stream, values, (int)values.Length, modulus, cache);
-        var mul = GetMulKernel(accelerator);
-        mul(stream, (int)values.Length, values, values, modulus);
+		ForwardDevice(accelerator, stream, values, (int)values.Length, modulus, cache);
+		var mulKernel = GetMulKernel(accelerator);
+        mulKernel(stream, (int)values.Length, values, values, modulus);
         InverseDevice(accelerator, stream, values, (int)values.Length, modulus, cache);
     }
 
     public static void PointwiseMultiply(Span<GpuUInt128> left, Span<GpuUInt128> right, GpuUInt128 modulus)
     {
-        var accelerator = SharedGpuContext.Accelerator;
+        var accelerator = AcceleratorPool.Shared.Rent();
         var stream = accelerator.CreateStream();
-        var kernel = GetMulKernel(accelerator);
         var pool = ThreadStaticPools.GpuUInt128Pool;
         var leftArray = pool.Rent(left.Length);
         var rightArray = pool.Rent(right.Length);
@@ -936,8 +1010,9 @@ public static class NttGpuMath
         var leftBuffer = accelerator.Allocate1D<GpuUInt128>(left.Length);
         var rightBuffer = accelerator.Allocate1D<GpuUInt128>(right.Length);
         leftBuffer.View.CopyFromCPU(stream, ref leftArray[0], left.Length);
-        rightBuffer.View.CopyFromCPU(stream, ref rightArray[0], right.Length);
-        kernel(stream, left.Length, leftBuffer.View, rightBuffer.View, modulus);
+		rightBuffer.View.CopyFromCPU(stream, ref rightArray[0], right.Length);
+		var mulKernel = GetMulKernel(accelerator);
+        mulKernel(stream, left.Length, leftBuffer.View, rightBuffer.View, modulus);
         leftBuffer.View.CopyToCPU(stream, ref leftArray[0], left.Length);
         rightBuffer.View.CopyToCPU(stream, ref rightArray[0], right.Length);
         stream.Synchronize();

PerfectNumbers.Core/Gpu/Pow2MontgomeryGpuCalculator.cs

@@ -0,0 +1,83 @@
+using System.Runtime.CompilerServices;
+using ILGPU;
+using ILGPU.Runtime;
+using PerfectNumbers.Core.Gpu.Accelerators;
+
+namespace PerfectNumbers.Core.Gpu;
+
+public static class Pow2MontgomeryGpuCalculator
+{
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static ulong CalculateKeep(ulong exponent, in MontgomeryDivisorData divisor)
+	{
+		GpuPrimeWorkLimiter.Acquire();
+
+		var gpu = Pow2MontgomeryAccelerator.Rent();
+		Accelerator? accelerator = gpu.Accelerator;
+		AcceleratorStream stream = gpu.Stream!;
+		MemoryBuffer1D<ulong, Stride1D.Dense> exponentBuffer = gpu.Input;
+		MemoryBuffer1D<ulong, Stride1D.Dense> resultBuffer = gpu.Output;
+
+		exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1);
+		// We don't need to worry about any left-overs here.
+		// resultBuffer.MemSetToZero(stream);
+
+		var kernel = gpu.KeepMontgomeryKernel;
+		kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View);
+
+		ulong result = 0UL;
+		resultBuffer.View.CopyToCPU(stream, ref result, 1);
+		stream.Synchronize();
+
+		Pow2MontgomeryAccelerator.Return(gpu);
+		GpuPrimeWorkLimiter.Release();
+		return result;
+	}
+
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static ulong CalculateConvert(ulong exponent, in MontgomeryDivisorData divisor)
+	{
+		GpuPrimeWorkLimiter.Acquire();
+
+		var gpu = Pow2MontgomeryAccelerator.Rent();
+		AcceleratorStream stream = gpu.Stream!;
+		MemoryBuffer1D<ulong, Stride1D.Dense> exponentBuffer = gpu.Input;
+		MemoryBuffer1D<ulong, Stride1D.Dense> resultBuffer = gpu.Output;
+
+		exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1);
+		// We don't need to worry about any left-overs here.
+		// resultBuffer.MemSetToZero(stream);
+
+		var kernel = gpu.ConvertToStandardKernel;
+		kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View);
+
+		ulong result = 0UL;
+		resultBuffer.View.CopyToCPU(stream, ref result, 1);
+		stream.Synchronize();
+
+		Pow2MontgomeryAccelerator.Return(gpu);
+		GpuPrimeWorkLimiter.Release();
+		return result;
+	}
+
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static ulong CalculateConvert(Pow2MontgomeryAccelerator gpu, ulong exponent, in MontgomeryDivisorData divisor)
+	{
+		AcceleratorStream stream = gpu.Stream!;
+		MemoryBuffer1D<ulong, Stride1D.Dense> exponentBuffer = gpu.Input;
+		MemoryBuffer1D<ulong, Stride1D.Dense> resultBuffer = gpu.Output;
+
+		exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1);
+		// We don't need to worry about any left-overs here.
+		// resultBuffer.MemSetToZero(stream);
+
+		var kernel = gpu.ConvertToStandardKernel;
+		kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View);
+
+		ulong result = 0UL;
+		resultBuffer.View.CopyToCPU(stream, ref result, 1);
+		stream.Synchronize();
+
+		return result;
+	}
+}

PerfectNumbers.Core/Gpu/SharedGpuContext.cs

@@ -1,5 +1,4 @@
-using System;
-using System.Threading;
+using System.Runtime.CompilerServices;
 using ILGPU;
 using ILGPU.Runtime;
 
@@ -7,14 +6,24 @@ namespace PerfectNumbers.Core.Gpu;
 
 public static class SharedGpuContext
 {
-	// TODO: Remove Lazy<T> in favor for better performance with nullable fields.
-    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);
-    private static readonly Lazy<Accelerator> s_accelerator = new(() => s_device.Value.CreateAccelerator(s_context.Value), LazyThreadSafetyMode.ExecutionAndPublication);
+	public static readonly Context Context = Context.CreateDefault();
+	public static readonly Device Device = Context.GetPreferredDevice(false);
 
-	internal static Context Context => s_context.Value;
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static Accelerator CreateAccelerator() => Device.CreateAccelerator(Context);
+}
 
-    internal static Device Device => s_device.Value;
+public readonly struct GpuContext
+{
+	public readonly Context Context;
+	public readonly Device Device;
 
-    internal static Accelerator Accelerator => s_accelerator.Value;
+	public GpuContext()
+	{
+		var context = Context.CreateDefault();
+		Context = context;
+		Device = context.GetPreferredDevice(false);
+	}
+
+	public Accelerator CreateAccelerator() => Device.CreateAccelerator(Context);
 }

PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs

@@ -2,7 +2,6 @@ using Open.Numeric.Primes;
 using ILGPU.Runtime;
 using System.Runtime.CompilerServices;
 using PerfectNumbers.Core.Gpu;
-using System.Runtime.InteropServices;
 
 namespace PerfectNumbers.Core;
 
@@ -41,7 +40,6 @@ public sealed class HeuristicCombinedPrimeTester
 		CombinedDivisorsEnding9Squares = CombinedDivisorsEnding9OneAOneBSquares;
 	}
 
-
 	public static void EnsureInitialized()
 	{
 		// Intentionally left blank. Accessing this method forces the static constructor to run.
@@ -229,7 +227,7 @@ public sealed class HeuristicCombinedPrimeTester
 		var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(1);
 		var accelerator = gpu.Accelerator;
 		var stream = gpu.Stream;
-		var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.HeuristicTrialDivisionKernel;
+		var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetHeuristicTrialDivisionKernel(accelerator);
 		var flagView1D = gpu.HeuristicFlag.View;
 		var flagView = flagView1D.AsContiguous();
 
@@ -267,109 +265,97 @@ public sealed class HeuristicCombinedPrimeTester
 		return n;
 	}
 
-	private static ReadOnlySpan<ulong> GetCombinedDivisors(byte nMod10)
+	private static ReadOnlySpan<ulong> GetCombinedDivisors(byte nMod10) => nMod10 switch
 	{
-		return nMod10 switch
+		1 => CombinedDivisorsEnding1,
+		3 => CombinedDivisorsEnding3,
+		7 => CombinedDivisorsEnding7,
+		9 => CombinedDivisorsEnding9,
+		_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
+	};
+
+	private static ReadOnlySpan<ulong> GetCombinedDivisorSquares(byte nMod10) => nMod10 switch
+	{
+		1 => CombinedDivisorsEnding1Squares,
+		3 => CombinedDivisorsEnding3Squares,
+		7 => CombinedDivisorsEnding7Squares,
+		9 => CombinedDivisorsEnding9Squares,
+		_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
+	};
+
+	internal static ReadOnlySpan<ulong> GetCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern) => pattern switch
+	{
+		CombinedDivisorPattern.TwoAOneB => nMod10 switch
 		{
-			1 => CombinedDivisorsEnding1,
-			3 => CombinedDivisorsEnding3,
-			7 => CombinedDivisorsEnding7,
-			9 => CombinedDivisorsEnding9,
+			1 => CombinedDivisorsEnding1TwoAOneB,
+			3 => CombinedDivisorsEnding3TwoAOneB,
+			7 => CombinedDivisorsEnding7TwoAOneB,
+			9 => CombinedDivisorsEnding9TwoAOneB,
 			_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
-		};
-	}
-
-	private static ReadOnlySpan<ulong> GetCombinedDivisorSquares(byte nMod10)
-	{
-		return nMod10 switch
+		},
+		CombinedDivisorPattern.OneAOneB => nMod10 switch
 		{
-			1 => CombinedDivisorsEnding1Squares,
-			3 => CombinedDivisorsEnding3Squares,
-			7 => CombinedDivisorsEnding7Squares,
-			9 => CombinedDivisorsEnding9Squares,
+			1 => CombinedDivisorsEnding1OneAOneB,
+			3 => CombinedDivisorsEnding3OneAOneB,
+			7 => CombinedDivisorsEnding7OneAOneB,
+			9 => CombinedDivisorsEnding9OneAOneB,
+			_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
+		},
+		CombinedDivisorPattern.ThreeAOneB => nMod10 switch
+		{
+			1 => CombinedDivisorsEnding1ThreeAOneB,
+			3 => CombinedDivisorsEnding3ThreeAOneB,
+			7 => CombinedDivisorsEnding7ThreeAOneB,
+			9 => CombinedDivisorsEnding9ThreeAOneB,
+			_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
+		},
+		CombinedDivisorPattern.ThreeATwoB => nMod10 switch
+		{
+			1 => CombinedDivisorsEnding1ThreeATwoB,
+			3 => CombinedDivisorsEnding3ThreeATwoB,
+			7 => CombinedDivisorsEnding7ThreeATwoB,
+			9 => CombinedDivisorsEnding9ThreeATwoB,
+			_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
+		},
+		_ => throw new InvalidOperationException($"Unsupported combined divisor pattern: {pattern}."),
+	};
+
+	internal static ReadOnlySpan<ulong> GetCombinedDivisorSquares(byte nMod10, CombinedDivisorPattern pattern) => pattern switch
+	{
+		CombinedDivisorPattern.TwoAOneB => nMod10 switch
+		{
+			1 => CombinedDivisorsEnding1TwoAOneBSquares,
+			3 => CombinedDivisorsEnding3TwoAOneBSquares,
+			7 => CombinedDivisorsEnding7TwoAOneBSquares,
+			9 => CombinedDivisorsEnding9TwoAOneBSquares,
 			_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
-		};
-	}
-
-	internal static ReadOnlySpan<ulong> GetCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern)
-	{
-		return pattern switch
+		},
+		CombinedDivisorPattern.OneAOneB => nMod10 switch
 		{
-			CombinedDivisorPattern.TwoAOneB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1TwoAOneB,
-				3 => CombinedDivisorsEnding3TwoAOneB,
-				7 => CombinedDivisorsEnding7TwoAOneB,
-				9 => CombinedDivisorsEnding9TwoAOneB,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
-			},
-			CombinedDivisorPattern.OneAOneB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1OneAOneB,
-				3 => CombinedDivisorsEnding3OneAOneB,
-				7 => CombinedDivisorsEnding7OneAOneB,
-				9 => CombinedDivisorsEnding9OneAOneB,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
-			},
-			CombinedDivisorPattern.ThreeAOneB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1ThreeAOneB,
-				3 => CombinedDivisorsEnding3ThreeAOneB,
-				7 => CombinedDivisorsEnding7ThreeAOneB,
-				9 => CombinedDivisorsEnding9ThreeAOneB,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
-			},
-			CombinedDivisorPattern.ThreeATwoB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1ThreeATwoB,
-				3 => CombinedDivisorsEnding3ThreeATwoB,
-				7 => CombinedDivisorsEnding7ThreeATwoB,
-				9 => CombinedDivisorsEnding9ThreeATwoB,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor selector for digit {nMod10}."),
-			},
-			_ => throw new InvalidOperationException($"Unsupported combined divisor pattern: {pattern}."),
-		};
-	}
-
-	internal static ReadOnlySpan<ulong> GetCombinedDivisorSquares(byte nMod10, CombinedDivisorPattern pattern)
-	{
-		return pattern switch
+			1 => CombinedDivisorsEnding1OneAOneBSquares,
+			3 => CombinedDivisorsEnding3OneAOneBSquares,
+			7 => CombinedDivisorsEnding7OneAOneBSquares,
+			9 => CombinedDivisorsEnding9OneAOneBSquares,
+			_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
+		},
+		CombinedDivisorPattern.ThreeAOneB => nMod10 switch
 		{
-			CombinedDivisorPattern.TwoAOneB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1TwoAOneBSquares,
-				3 => CombinedDivisorsEnding3TwoAOneBSquares,
-				7 => CombinedDivisorsEnding7TwoAOneBSquares,
-				9 => CombinedDivisorsEnding9TwoAOneBSquares,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
-			},
-			CombinedDivisorPattern.OneAOneB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1OneAOneBSquares,
-				3 => CombinedDivisorsEnding3OneAOneBSquares,
-				7 => CombinedDivisorsEnding7OneAOneBSquares,
-				9 => CombinedDivisorsEnding9OneAOneBSquares,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
-			},
-			CombinedDivisorPattern.ThreeAOneB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1ThreeAOneBSquares,
-				3 => CombinedDivisorsEnding3ThreeAOneBSquares,
-				7 => CombinedDivisorsEnding7ThreeAOneBSquares,
-				9 => CombinedDivisorsEnding9ThreeAOneBSquares,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
-			},
-			CombinedDivisorPattern.ThreeATwoB => nMod10 switch
-			{
-				1 => CombinedDivisorsEnding1ThreeATwoBSquares,
-				3 => CombinedDivisorsEnding3ThreeATwoBSquares,
-				7 => CombinedDivisorsEnding7ThreeATwoBSquares,
-				9 => CombinedDivisorsEnding9ThreeATwoBSquares,
-				_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
-			},
-			_ => throw new InvalidOperationException($"Unsupported combined divisor pattern for squares: {pattern}."),
-		};
-	}
+			1 => CombinedDivisorsEnding1ThreeAOneBSquares,
+			3 => CombinedDivisorsEnding3ThreeAOneBSquares,
+			7 => CombinedDivisorsEnding7ThreeAOneBSquares,
+			9 => CombinedDivisorsEnding9ThreeAOneBSquares,
+			_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
+		},
+		CombinedDivisorPattern.ThreeATwoB => nMod10 switch
+		{
+			1 => CombinedDivisorsEnding1ThreeATwoBSquares,
+			3 => CombinedDivisorsEnding3ThreeATwoBSquares,
+			7 => CombinedDivisorsEnding7ThreeATwoBSquares,
+			9 => CombinedDivisorsEnding9ThreeATwoBSquares,
+			_ => throw new InvalidOperationException($"Unsupported combined divisor square selector for digit {nMod10}."),
+		},
+		_ => throw new InvalidOperationException($"Unsupported combined divisor pattern for squares: {pattern}."),
+	};
 
 	private static ulong[] BuildCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern, out ulong[] squares)
 	{
@@ -443,8 +429,6 @@ public sealed class HeuristicCombinedPrimeTester
 		return result;
 	}
 
-
-
 	private static int GetGroupBStartIndex(ReadOnlySpan<uint> divisors)
 	{
 		int index = 0;

PerfectNumbers.Core/HeuristicPrimeTester.cs

@@ -244,7 +244,7 @@ public sealed class HeuristicPrimeTester
 		var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(1);
 		var accelerator = gpu.Accelerator;
 		var stream = gpu.Stream;
-		var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.HeuristicTrialDivisionKernel;
+		var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetHeuristicTrialDivisionKernel(accelerator);
 		var flagView1D = gpu.HeuristicFlag.View;
 		var flagView = flagView1D.AsContiguous();
 

PerfectNumbers.Core/MersenneDivisorCycles.cs

@@ -48,9 +48,7 @@ public class MersenneDivisorCycles
 	public ulong[] ExportSmallCyclesSnapshot()
 	{
 		EnsureSmallBuffer();
-		ulong[] snapshot = new ulong[PerfectNumberConstants.MaxQForDivisorCycles + 1];
-		Array.Copy(_smallCycles!, snapshot, snapshot.Length);
-		return snapshot;
+		return _smallCycles!;
 	}
 
 	public static bool CycleEqualsExponent(ulong divisor, in MontgomeryDivisorData divisorData, ulong exponent)
@@ -790,6 +788,9 @@ public class MersenneDivisorCycles
 		Task.WaitAll(tasks);
 	}
 
+	[ThreadStatic]
+	private static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>>? _divisorCycleKernel;
+
 	public static void GenerateGpu(string path, ulong maxDivisor, int batchSize = 1_000_000, long skipCount = 0L, long nextPosition = 0L)
 	{
 		// Prepare output file with header
@@ -804,15 +805,8 @@ public class MersenneDivisorCycles
 		divisors = pool.Rent(count);
 		outCycles = pool.Rent(count);
 		GpuPrimeWorkLimiter.Acquire();
-		var accelerator = SharedGpuContext.Accelerator;
+		var accelerator = AcceleratorPool.Shared.Rent();
 		var stream = accelerator.CreateStream();
-		var loaded = accelerator.LoadAutoGroupedStreamKernel<
-				Index1D,
-				ArrayView1D<ulong, Stride1D.Dense>,
-				ArrayView1D<ulong, Stride1D.Dense>>(DivisorCycleKernels.GpuDivisorCycleKernel);
-		var kernel = KernelUtil.GetKernel(loaded);
-		var launcher = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>>();
-
 
 		using Stream outputStream = new FileStream(path, FileMode.OpenOrCreate, FileAccess.ReadWrite, FileShare.Read, BufferSize10M, useAsync: true);
 		if (nextPosition > 0L)
@@ -841,6 +835,7 @@ public class MersenneDivisorCycles
 			skipCount--;
 		}
 
+		var divisorCycleKernel = GetDivisorCycleKernel(accelerator);
 		while (start <= maxDivisor)
 		{
 			end = Math.Min(start + batchSizeUL - 1UL, maxDivisor);
@@ -859,10 +854,15 @@ public class MersenneDivisorCycles
 
 			// Use GpuKernelPool to get a kernel and context
 			// GpuKernelPool.Run((accelerator, stream) =>
-			var bufferDiv = accelerator.Allocate1D(validDivisors);
-			var bufferCycle = accelerator.Allocate1D<ulong>(idx);
+			MemoryBuffer1D<ulong, Stride1D.Dense>? bufferDiv;
+			MemoryBuffer1D<ulong, Stride1D.Dense>? bufferCycle;
+			// lock(accelerator)
+			{
+				bufferDiv = accelerator.Allocate1D(validDivisors);
+				bufferCycle = accelerator.Allocate1D<ulong>(idx);				
+			}
 
-			launcher(
+			divisorCycleKernel(
 					stream,
 					idx,
 					bufferDiv.View,
@@ -900,6 +900,25 @@ public class MersenneDivisorCycles
 		GpuPrimeWorkLimiter.Release();
 	}
 
+	private static Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> GetDivisorCycleKernel(Accelerator accelerator)
+	{
+		var pool = _divisorCycleKernel ??= [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		var loaded = accelerator.LoadAutoGroupedStreamKernel<
+				Index1D,
+				ArrayView1D<ulong, Stride1D.Dense>,
+				ArrayView1D<ulong, Stride1D.Dense>>(DivisorCycleKernels.GpuDivisorCycleKernel);
+
+		var kernel = KernelUtil.GetKernel(loaded);
+
+		cached = kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
 
 	public static ulong CalculateCycleLengthGpu(ulong divisor) => DivisorCycleKernels.CalculateCycleLengthGpu(divisor);
 

PerfectNumbers.Core/MersenneNumberDivisorByDivisorTester.cs

@@ -293,7 +293,7 @@ public static class MersenneNumberDivisorByDivisorTester
 				tasks[taskIndex++] = Task.Factory.StartNew(
 					() =>
 					{
-						startGate.Wait();
+						// startGate.Wait();
 						Console.WriteLine($"Task started for range {rangeStart}");
 
 						for (int index = rangeStart; index < rangeEnd; index++)

PerfectNumbers.Core/MersenneNumberTester.cs

@@ -33,7 +33,7 @@ public sealed class MersenneNumberTester(
 	private readonly bool _useGpuLucas = useGpuLucas;
 	private readonly bool _useGpuScan = useGpuScan;     // device for pow2mod/incremental scanning
 	private readonly bool _useGpuOrder = useGpuOrder;   // device for order computations
-	private static readonly ConcurrentDictionary<UInt128, ulong> OrderCache = new();
+	private static readonly ConcurrentDictionary<UInt128, ulong> _orderCache = [];
 	// TODO: Swap this ConcurrentDictionary for the pooled dictionary variant highlighted in
 	// Pow2MontgomeryModBenchmarks once order warmups reuse deterministic divisor-cycle snapshots; the pooled approach removed
 	// the locking overhead when scanning p >= 138M in those benchmarks.
@@ -89,13 +89,13 @@ public sealed class MersenneNumberTester(
 					}
 				}
 
-				if (!shouldCheck || !cacheEnabled || (cacheEnabled && OrderCache.ContainsKey(q)))
+				if (!shouldCheck || !cacheEnabled || (cacheEnabled && _orderCache.ContainsKey(q)))
 				{
 					continue;
 				}
 
 				ulong ord = q.CalculateOrder();
-				OrderCache[q] = ord;
+				_orderCache[q] = ord;
 			}
 
 			return;
@@ -131,7 +131,7 @@ public sealed class MersenneNumberTester(
 				}
 			}
 
-			if (!shouldCheck2 || (cacheEnabled && OrderCache.ContainsKey(q2)))
+			if (!shouldCheck2 || (cacheEnabled && _orderCache.ContainsKey(q2)))
 			{
 				continue;
 			}
@@ -148,9 +148,9 @@ public sealed class MersenneNumberTester(
 		UInt128[] qs = qsBuffer;
 
 		GpuPrimeWorkLimiter.Acquire();
-		var accelerator = SharedGpuContext.Accelerator;
+		var accelerator = AcceleratorPool.Shared.Rent();
 		var stream = accelerator.CreateStream();
-		var orderKernel = GpuKernelPool.Kernels.Order!;
+		var orderKernel = GpuKernelPool.GetOrAddKernels(accelerator, stream).Order!;
 
 		// Guard long-running kernels by chunking the warm-up across batches.
 		// Reuse the same ScanBatchSize knob used for scanning.
@@ -158,11 +158,18 @@ public sealed class MersenneNumberTester(
 		ulong divMul = (ulong)((((UInt128)1 << 64) - 1UL) / exponent) + 1UL;
 		int offset = 0;
 		var gpuUInt128Pool = ThreadStaticPools.GpuUInt128Pool;
+		MemoryBuffer1D<GpuUInt128, Stride1D.Dense>? qBuffer;
+		MemoryBuffer1D<ulong, Stride1D.Dense>? orderBuffer;
+		
 		while (offset < idx)
 		{
 			int count = Math.Min(batchSize, idx - offset);
-			var qBuffer = accelerator.Allocate1D<GpuUInt128>(count);
-			var orderBuffer = accelerator.Allocate1D<ulong>(count);
+			// lock (accelerator)
+			{
+				qBuffer = accelerator.Allocate1D<GpuUInt128>(count);
+				orderBuffer = accelerator.Allocate1D<ulong>(count);
+			}
+			
 			// Convert to device-friendly GpuUInt128 on the fly
 			var tmp = gpuUInt128Pool.Rent(count);
 			for (int i = 0; i < count; i++)
@@ -187,7 +194,7 @@ public sealed class MersenneNumberTester(
 
 				if (useOrderCache)
 				{
-					OrderCache[qs[offset + i]] = order;
+					_orderCache[qs[offset + i]] = order;
 				}
 			}
 

PerfectNumbers.Core/PerfectNumberConstants.cs

@@ -1,3 +1,5 @@
+using PerfectNumbers.Core.Gpu;
+
 namespace PerfectNumbers.Core;
 
 public static class PerfectNumberConstants
@@ -8,7 +10,8 @@ public static class PerfectNumberConstants
     public const int MaxQForDivisorCycles = 4_000_000;
     public static readonly uint PrimesLimit = 1_000_000; //(ulong)Array.MaxLength;// 1_000_000;
     public const int MaxOddPowersCount = 128;
-    public const int PooledArrayThreshold = 64;
+	public const int PooledArrayThreshold = 64;
+	public static readonly int RollingAccelerators = 128; //SharedGpuContext.Device.MaxNumThreadsPerGroup;
     // TODO: Load these limits from the benchmark-driven configuration so CPU and GPU scans stay aligned with the optimal
     // divisor-cycle datasets we generate offline.
     // TODO: Promote these magic numbers into a runtime profile derived from EvenPerfectBitScanner.Benchmarks so we can retune

PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs

@@ -3,8 +3,8 @@ using System.Diagnostics;
 using System.Numerics;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
-using ILGPU;
 using PerfectNumbers.Core.Gpu;
+using PerfectNumbers.Core.Gpu.Accelerators;
 
 namespace PerfectNumbers.Core;
 
@@ -52,8 +52,10 @@ internal static partial class PrimeOrderCalculator
 
 		ulong phi = prime - 1UL;
 
+		Pow2MontgomeryAccelerator gpu = Pow2MontgomeryAccelerator.Rent();
 		if (IsGpuHeuristicDevice && PrimeOrderGpuHeuristics.TryCalculateOrder(prime, previousOrder, config, divisorData, out ulong gpuOrder))
 		{
+			Pow2MontgomeryAccelerator.Return(gpu);
 			return gpuOrder;
 		}
 
@@ -62,17 +64,19 @@ internal static partial class PrimeOrderCalculator
 		ulong result;
 		if (phiFactors.Factors is null)
 		{
-			result = CalculateByFactorizationCpu(prime, divisorData);
+			result = CalculateByFactorizationCpu(gpu, prime, divisorData);
 			phiFactors.Dispose();
 			return result;
 		}
 
-		result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors);
+		result = RunHeuristicPipelineCpu(gpu, prime, previousOrder, config, divisorData, phi, phiFactors);
 		phiFactors.Dispose();
+		Pow2MontgomeryAccelerator.Return(gpu);
 		return result;
 	}
 
 	private static ulong RunHeuristicPipelineCpu(
+		Pow2MontgomeryAccelerator gpu,
 		ulong prime,
 		ulong? previousOrder,
 		in PrimeOrderSearchConfig config,
@@ -85,25 +89,30 @@ internal static partial class PrimeOrderCalculator
 			return phi;
 		}
 
-		ulong candidateOrder = InitializeStartingOrderCpu(prime, phi, divisorData);
+		GpuPrimeWorkLimiter.Acquire();
+		ulong candidateOrder = InitializeStartingOrderCpu(gpu, prime, phi, divisorData);
 		candidateOrder = ExponentLoweringCpu(candidateOrder, prime, phiFactors, divisorData);
 
-		if (TryConfirmOrderCpu(prime, candidateOrder, divisorData, config, out PartialFactorResult? orderFactors))
+		if (TryConfirmOrderCpu(gpu, prime, candidateOrder, divisorData, config, out PartialFactorResult? orderFactors))
 		{
+			GpuPrimeWorkLimiter.Release();
 			return candidateOrder;
 		}
 
 		if (config.Mode == PrimeOrderMode.Strict)
 		{
 			orderFactors?.Dispose();
-			return CalculateByFactorizationCpu(prime, divisorData);
+			GpuPrimeWorkLimiter.Release();
+			return CalculateByFactorizationCpu(gpu, prime, divisorData);
 		}
 
 		if (TryHeuristicFinishCpu(prime, candidateOrder, previousOrder, divisorData, config, phiFactors, orderFactors, out ulong order))
 		{
+			GpuPrimeWorkLimiter.Release();
 			return order;
 		}
 
+		GpuPrimeWorkLimiter.Release();
 		return candidateOrder;
 	}
 
@@ -185,13 +194,13 @@ internal static partial class PrimeOrderCalculator
 		return true;
 	}
 
-	private static ulong InitializeStartingOrderCpu(ulong prime, ulong phi, in MontgomeryDivisorData divisorData)
+	private static ulong InitializeStartingOrderCpu(Pow2MontgomeryAccelerator gpu, ulong prime, ulong phi, in MontgomeryDivisorData divisorData)
 	{
 		ulong order = phi;
 		if ((prime & 7UL) == 1UL || (prime & 7UL) == 7UL)
 		{
 			ulong half = phi >> 1;
-			if (Pow2EqualsOneCpu(half, prime, divisorData))
+			if (Pow2EqualsOneCpu(gpu, half, prime, divisorData))
 			{
 				order = half;
 			}
@@ -319,6 +328,7 @@ internal static partial class PrimeOrderCalculator
 	}
 
 	private static bool TryConfirmOrderCpu(
+		Pow2MontgomeryAccelerator gpu,
 		ulong prime,
 		ulong order,
 		in MontgomeryDivisorData divisorData,
@@ -332,7 +342,7 @@ internal static partial class PrimeOrderCalculator
 			return false;
 		}
 
-		if (!Pow2EqualsOneCpu(order, prime, divisorData))
+		if (!Pow2EqualsOneCpu(gpu, order, prime, divisorData))
 		{
 			return false;
 		}
@@ -1021,11 +1031,15 @@ internal static partial class PrimeOrderCalculator
 		return false;
 	}
 
-	private static bool Pow2EqualsOneCpu(ulong exponent, ulong prime, in MontgomeryDivisorData divisorData)
+	private static ulong _pow2GpuHits;
+	private static ulong _pow2CpuHits;
+	private static bool Pow2EqualsOneCpu(Pow2MontgomeryAccelerator gpu, ulong exponent, ulong prime, in MontgomeryDivisorData divisorData)
 	{
-		if (IsGpuPow2Allowed)
+		if (exponent <= (PerfectNumberConstants.MaxQForDivisorCycles) || prime <= (PerfectNumberConstants.MaxQForDivisorCycles))
 		{
-			ulong remainder = exponent.Pow2MontgomeryModWindowedGpu(divisorData, false);
+			// Atomic.Add(ref _pow2GpuHits, 1);
+			// Console.WriteLine($"pow2 GPU Hits {_pow2GpuHits}");
+			ulong remainder = exponent.Pow2MontgomeryModWindowedConvertGpu(gpu, divisorData);
 			// GpuPow2ModStatus status = PrimeOrderGpuHeuristics.TryPow2Mod(exponent, prime, out ulong remainder, divisorData);
 			// if (status == GpuPow2ModStatus.Success)
 			// {
@@ -1033,6 +1047,8 @@ internal static partial class PrimeOrderCalculator
 			// }
 		}
 
+		// Atomic.Add(ref _pow2CpuHits, 1);
+		// Console.WriteLine($"pow2 CPU Hits {_pow2CpuHits}");
 		return exponent.Pow2MontgomeryModWindowedCpu(divisorData, keepMontgomery: false) == 1UL;
 	}
 
@@ -1638,7 +1654,7 @@ internal static partial class PrimeOrderCalculator
 		exponents = newExponentSpan;
 	}
 
-	private static ulong CalculateByFactorizationCpu(ulong prime, in MontgomeryDivisorData divisorData)
+	private static ulong CalculateByFactorizationCpu(Pow2MontgomeryAccelerator gpu, ulong prime, in MontgomeryDivisorData divisorData)
 	{
 		ulong phi = prime - 1UL;
 		Dictionary<ulong, int> counts = new(capacity: 8);
@@ -1653,6 +1669,8 @@ internal static partial class PrimeOrderCalculator
 
 		ulong order = phi;
 		int entryCount = entries.Count;
+
+		GpuPrimeWorkLimiter.Acquire();
 		for (int i = 0; i < entryCount; i++)
 		{
 			ulong primeFactor = entries[i].Key;
@@ -1665,7 +1683,10 @@ internal static partial class PrimeOrderCalculator
 				}
 
 				ulong candidate = order / primeFactor;
-				if (Pow2EqualsOneCpu(candidate, prime, divisorData))
+
+				// if (candidate.Pow2MontgomeryModWindowedCpu(divisorData, keepMontgomery: false) == 1UL)
+				if (candidate.Pow2MontgomeryModWindowedConvertGpu(gpu, divisorData) == 1UL)
+				// if (Pow2EqualsOneCpu(candidate, prime, divisorData))
 				{
 					order = candidate;
 					continue;
@@ -1675,6 +1696,7 @@ internal static partial class PrimeOrderCalculator
 			}
 		}
 
+		GpuPrimeWorkLimiter.Release();
 		return order;
 	}
 

PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs

@@ -1,6 +1,7 @@
 using System.Runtime.InteropServices;
 using ILGPU;
 using ILGPU.Runtime;
+using ILGPU.Runtime.Cuda;
 using PerfectNumbers.Core.Gpu;
 
 namespace PerfectNumbers.Core;
@@ -20,18 +21,18 @@ internal static partial class PrimeOrderCalculator
 		remaining = value;
 
 		GpuPrimeWorkLimiter.Acquire();
-		Accelerator accelerator = SharedGpuContext.Accelerator;
+		Accelerator accelerator = AcceleratorPool.Shared.Rent();
 		var stream = accelerator.CreateStream();
-		KernelContainer kernels = GpuKernelPool.Kernels;
+		KernelContainer kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream);
 
 		ScratchBuffer scratch = GpuScratchBufferPool.Rent(accelerator, GpuSmallPrimeFactorSlots, 0);
 
-		SmallPrimeFactorViews tables = GpuKernelPool.GetSmallPrimeFactorTables();
+		SmallPrimeFactorViews tables = GpuKernelPool.GetSmallPrimeFactorTables(kernels);
 
 		if (kernels.SmallPrimeFactorsPrimes is null) throw new NullReferenceException($"{nameof(kernels.SmallPrimeFactorsPrimes)} is null");
 		if (kernels.SmallPrimeFactorsSquares is null) throw new NullReferenceException($"{nameof(kernels.SmallPrimeFactorsSquares)} is null");
 
-		var kernel = GpuKernelPool.Kernels.SmallPrimeFactor!;
+		var kernel = kernels.SmallPrimeFactor!;
 
 		ArrayView1D<int, Stride1D.Dense> smallPrimeFactorCountSlotView = scratch.SmallPrimeFactorCountSlot.View;
 		ArrayView1D<ulong, Stride1D.Dense> smallPrimeFactorRemainingSlotView = scratch.SmallPrimeFactorRemainingSlot.View;
@@ -53,10 +54,11 @@ internal static partial class PrimeOrderCalculator
 				smallPrimeFactorRemainingSlotView);
 
 
+		// Span<int> factorCountTemp = stackalloc int[1];
 		factorCount = 0;
 		smallPrimeFactorCountSlotView.CopyToCPU(stream, ref factorCount, 1);
 		stream.Synchronize();
-
+		// factorCount = factorCountTemp[0];
 		factorCount = Math.Min(factorCount, GpuSmallPrimeFactorSlots);
 
 		if (factorCount > 0)
@@ -66,7 +68,10 @@ internal static partial class PrimeOrderCalculator
 			smallPrimeFactorExponentSlotsView.CopyToCPU(stream, ref MemoryMarshal.GetReference(exponentBuffer), factorCount);
 		}
 
+		// Span<ulong> remainingTemp = stackalloc ulong[1];
+
 		smallPrimeFactorRemainingSlotView.CopyToCPU(stream, ref remaining, 1);
+		// remaining = remainingTemp[0];
 		stream.Synchronize();
 		stream.Dispose();
 
@@ -102,16 +107,16 @@ internal static partial class PrimeOrderCalculator
 		}
 
 		GpuPrimeWorkLimiter.Acquire();
-		Accelerator accelerator = SharedGpuContext.Accelerator;
-
+		Accelerator accelerator = AcceleratorPool.Shared.Rent();
 		var stream = accelerator.CreateStream();
+		var kernels = GpuKernelPool.GetOrAddKernels(accelerator, stream);
 
 		ScratchBuffer scratch = GpuScratchBufferPool.Rent(accelerator, 0, factorCount);
 
 		ArrayView1D<ulong, Stride1D.Dense> specialMaxFactorsView = scratch.SpecialMaxFactors.View;
 		specialMaxFactorsView.SubView(0, factorCount).CopyFromCPU(stream, ref MemoryMarshal.GetReference(factorSpan), factorCount);
 
-		var kernel = GpuKernelPool.Kernels.SpecialMax!;
+		var kernel = kernels.SpecialMax!;
 
 		ArrayView1D<ulong, Stride1D.Dense> specialMaxResultView = scratch.SpecialMaxResult.View;
 
@@ -125,9 +130,10 @@ internal static partial class PrimeOrderCalculator
 				scratch.SpecialMaxCandidates.View,
 				specialMaxResultView);
 
-		Span<ulong> result = stackalloc ulong[1];
+		// Span<ulong> result = stackalloc ulong[1];
+		ulong result = 0UL;
 		// specialMaxResultView.CopyToCPU(stream, ref result, 1);
-		specialMaxResultView.CopyToCPU(stream, in result);
+		specialMaxResultView.CopyToCPU(stream, ref result, 1);
 		stream.Synchronize();
 		stream.Dispose();
 
@@ -135,7 +141,7 @@ internal static partial class PrimeOrderCalculator
 		// return result != 0;
 		GpuScratchBufferPool.Return(scratch);
 		GpuPrimeWorkLimiter.Release();
-		return result[0] != 0;
+		return result != 0;
 	}
 }
 

PerfectNumbers.Core/PrimeTester.cs

@@ -99,9 +99,10 @@ public sealed class PrimeTester
 		var gpu = PrimeTesterGpuContextPool.Rent(1);
 		var accelerator = gpu.Accelerator;
 		var stream = gpu.Stream;
-		var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.Kernel;
+		var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetKernel(accelerator);
 
 		ulong value = n;
+		// Span<byte> flag = stackalloc byte[1];
 		byte flag = 0;
 
 		var input = gpu.Input;
@@ -166,7 +167,7 @@ public sealed class PrimeTester
 		var gpu = PrimeTesterGpuContextPool.Rent(GpuBatchSize);
 		var accelerator = gpu.Accelerator;
 		var stream = gpu.Stream;
-		var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.Kernel;
+		var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.GetKernel(accelerator);
 		int totalLength = values.Length;
 		int batchSize = GpuBatchSize;
 
@@ -215,83 +216,52 @@ public sealed class PrimeTester
 
 	internal static class PrimeTesterGpuContextPool
 	{
-		private static readonly ConcurrentQueue<PrimeTesterGpuContextLease> Pool = new();
+		[ThreadStatic]
+		private static Queue<PrimeTesterGpuContextLease>? Pool;
 		private static readonly object WarmUpLock = new();
 		private static int WarmedLeaseCount;
-		private static readonly object SharedTablesLock = new();
-		private static readonly Dictionary<Accelerator, SharedHeuristicGpuTables> SharedTables = new(AcceleratorReferenceComparer.Instance);
 
-		private static SharedHeuristicGpuTables RentSharedTables(Accelerator accelerator, AcceleratorStream stream)
-		{
-			lock (SharedTablesLock)
-			{
-				if (!SharedTables.TryGetValue(accelerator, out var tables))
-				{
-					tables = new SharedHeuristicGpuTables(accelerator, stream);
-					SharedTables.Add(accelerator, tables);
-				}
-
-				tables.AddReference();
-				return tables;
-			}
-		}
-
-		private static void ReleaseSharedTables(SharedHeuristicGpuTables tables)
-		{
-			if (tables is null)
-			{
-				return;
-			}
-
-			lock (SharedTablesLock)
-			{
-				tables.Release();
-			}
-		}
-
-		internal static void EnsureStaticTables(KernelContainer kernels, Accelerator accelerator, AcceleratorStream stream)
-		{
-			var tables = RentSharedTables(accelerator, stream);
-			var heuristicTables = tables.CreateHeuristicDivisorTables();
-			HeuristicGpuDivisorTables.InitializeShared(in heuristicTables);
-			ReleaseSharedTables(tables);
-		}
+		// internal static void EnsureStaticTables(Accelerator accelerator)
+		// {
+		// 	var tables = CreateSharedTables(accelerator);
+		// }
 
 		internal static void WarmUp(int target, int minBufferCapacity)
 		{
-			if (target <= 0)
-			{
-				return;
-			}
+			// if (target <= 0)
+			// {
+			// 	return;
+			// }
 
-			lock (WarmUpLock)
-			{
-				if (target <= WarmedLeaseCount)
-				{
-					return;
-				}
+			// lock (WarmUpLock)
+			// {
+			// 	if (target <= WarmedLeaseCount)
+			// 	{
+			// 		return;
+			// 	}
 
-				Accelerator accelerator = SharedGpuContext.Accelerator;
-				AcceleratorStream stream = accelerator.CreateStream();
-				KernelContainer kernels = GpuKernelPool.Kernels;
+			// 	Accelerator accelerator = SharedGpuContext.Accelerator;
+			// 	AcceleratorStream stream = accelerator.CreateStream();
+			// 	KernelContainer kernels = GpuKernelPool.Kernels;
 
-				GpuStaticTableInitializer.EnsureStaticTables(kernels, accelerator, stream);
+			// 	GpuStaticTableInitializer.EnsureStaticTables(kernels, stream);
 
-				int toCreate = target - WarmedLeaseCount;
-				for (int i = 0; i < toCreate; i++)
-				{
-					Pool.Enqueue(new PrimeTesterGpuContextLease(minBufferCapacity, stream));
-				}
+			// 	int toCreate = target - WarmedLeaseCount;
+			// 	for (int i = 0; i < toCreate; i++)
+			// 	{
+			// 		Pool.Enqueue(new PrimeTesterGpuContextLease(minBufferCapacity, stream));
+			// 	}
 
-				stream.Synchronize();
-				stream.Dispose();
-				WarmedLeaseCount = target;
-			}
+			// 	stream.Synchronize();
+			// 	stream.Dispose();
+			// 	WarmedLeaseCount = target;
+			// }
 		}
 
 		internal static PrimeTesterGpuContextLease Rent(int minBufferCapacity = 1)
 		{
-			if (!Pool.TryDequeue(out var lease))
+			var pool = Pool ??= new();
+			if (!pool.TryDequeue(out var lease))
 			{
 				lease = new PrimeTesterGpuContextLease(minBufferCapacity);
 			}
@@ -306,74 +276,37 @@ public sealed class PrimeTester
 
 		internal static void Return(PrimeTesterGpuContextLease lease)
 		{
-			Pool.Enqueue(lease);
+			var pool = Pool ??= new();
+				pool.Enqueue(lease);
 		}
 
 		internal static void DisposeAll()
 		{
-			while (Pool.TryDequeue(out var lease))
+			if (Pool != null)
 			{
-				lease.DisposeResources();
-			}
-
-			lock (SharedTablesLock)
-			{
-				foreach (var entry in SharedTables.Values)
+				while (Pool.TryDequeue(out var lease))
 				{
-					entry.Dispose();
+					lease.DisposeResources();
 				}
-
-				SharedTables.Clear();
-				var emptyTables = default(HeuristicGpuDivisorTables);
-				HeuristicGpuDivisorTables.InitializeShared(in emptyTables);
 			}
 
 			WarmedLeaseCount = 0;
 		}
 
-		internal static void Clear(Accelerator accelerator)
+		internal static void Clear()
 		{
-			if (accelerator is null)
+			if (Pool != null)
 			{
-				return;
-			}
-
-			lock (WarmUpLock)
-			{
-				var retained = new List<PrimeTesterGpuContextLease>();
-				while (Pool.TryDequeue(out var lease))
+				lock (WarmUpLock)
 				{
-					if (lease.Accelerator == accelerator)
+					var retained = new List<PrimeTesterGpuContextLease>();
+					while (Pool.TryDequeue(out var lease))
 					{
 						lease.DisposeResources();
 					}
-					else
-					{
-						retained.Add(lease);
-					}
-				}
 
-				int retainedCount = retained.Count;
-				for (int i = 0; i < retainedCount; i++)
-				{
-					Pool.Enqueue(retained[i]);
-				}
-
-				WarmedLeaseCount = Pool.Count;
-			}
-
-			lock (SharedTablesLock)
-			{
-				if (SharedTables.Remove(accelerator, out var tables))
-				{
-					tables.Dispose();
-				}
-
-				if (SharedTables.Count == 0)
-				{
-					var emptyTables = default(HeuristicGpuDivisorTables);
-					HeuristicGpuDivisorTables.InitializeShared(in emptyTables);
-				}
+					WarmedLeaseCount = Pool.Count;
+				}				
 			}
 		}
 
@@ -390,9 +323,38 @@ public sealed class PrimeTester
 				get => _stream ?? Accelerator.CreateStream();
 			}
 
-			public static readonly Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> Kernel = KernelUtil.GetKernel(SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel)).CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>>();
+			[ThreadStatic]
+			public static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>>? Kernel;
+
+			internal static Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> GetKernel(Accelerator accelerator)
+			{
+				var pool = Kernel ??= [];
+				if (pool.TryGetValue(accelerator, out var cached))
+				{
+					return cached;
+				}
+
+				cached = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel)).CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>>();
+				pool[accelerator] = cached;
+				return cached;
+			}
+
+			[ThreadStatic]
+			public static Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>>? HeuristicTrialDivisionKernel;
+
+			internal static Action<AcceleratorStream, Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables> GetHeuristicTrialDivisionKernel(Accelerator accelerator)
+			{
+				var pool = HeuristicTrialDivisionKernel ??= [];
+				if (pool.TryGetValue(accelerator, out var cached))
+				{
+					return cached;
+				}
+
+				cached = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel)).CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>>();
+				pool[accelerator] = cached;
+				return cached;
+			}
 
-			public static readonly Action<AcceleratorStream, Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables> HeuristicTrialDivisionKernel = KernelUtil.GetKernel(SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel)).CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>>();
 
 			public MemoryBuffer1D<ulong, Stride1D.Dense> Input = null!;
 			public MemoryBuffer1D<byte, Stride1D.Dense> Output = null!;
@@ -431,37 +393,27 @@ public sealed class PrimeTester
 			internal PrimeTesterGpuContextLease(int minBufferCapacity)
 			{
 				AcceleratorContext = SharedGpuContext.Context;
-				Accelerator = SharedGpuContext.Accelerator;
-				AcceleratorStream stream = SharedGpuContext.Accelerator.CreateStream();
-				var kernels = GpuKernelPool.Kernels;
-				GpuStaticTableInitializer.EnsureStaticTables(kernels, Accelerator, stream);
+				Accelerator accelerator = AcceleratorPool.Shared.Rent();
+				// lock (_accelerators)
+				// {
+				// }
 
-				_sharedTables = RentSharedTables(Accelerator, stream);
+				Accelerator = accelerator;
+				AcceleratorStream stream = accelerator.CreateStream();
+				// var kernels = _kernels[index]; //GpuKernelPool.Kernels;
+				// GpuStaticTableInitializer.EnsureStaticTables(accelerator, kernels, stream);
 
-				_heuristicDivisorTables = _sharedTables.CreateHeuristicDivisorTables();
-				HeuristicGpuDivisorTables.InitializeShared(in _heuristicDivisorTables);
+				var sharedTables = SharedHeuristicGpuTables.EnsureStaticTables(accelerator, stream);
+				_sharedTables = sharedTables;
+
+				var heuristicDivisorTables = sharedTables.CreateHeuristicDivisorTables();
+				_heuristicDivisorTables = heuristicDivisorTables;
 
 				EnsureCapacity(minBufferCapacity);
 				stream.Synchronize();
 				stream.Dispose();
 			}
 
-			internal PrimeTesterGpuContextLease(int minBufferCapacity, AcceleratorStream stream)
-			{
-				AcceleratorContext = SharedGpuContext.Context;
-				Accelerator = SharedGpuContext.Accelerator;
-				// GpuStaticTableInitializer.EnsureStaticTables(Accelerator);
-				// Kernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
-				// HeuristicTrialDivisionKernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel);
-
-				_sharedTables = RentSharedTables(Accelerator, stream);
-
-				_heuristicDivisorTables = _sharedTables.CreateHeuristicDivisorTables();
-				HeuristicGpuDivisorTables.InitializeShared(in _heuristicDivisorTables);
-
-				EnsureCapacity(minBufferCapacity);
-			}
-
 			public void EnsureCapacity(int minCapacity)
 			{
 				if (minCapacity <= BufferCapacity)
@@ -473,14 +425,20 @@ public sealed class PrimeTester
 				Output?.Dispose();
 				HeuristicFlag?.Dispose();
 
-				Input = Accelerator.Allocate1D<ulong>(minCapacity);
-				Output = Accelerator.Allocate1D<byte>(minCapacity);
-				HeuristicFlag = Accelerator.Allocate1D<int>(minCapacity);
+				// lock (Accelerator)
+				{
+					Input = Accelerator.Allocate1D<ulong>(minCapacity);
+					Output = Accelerator.Allocate1D<byte>(minCapacity);
+					HeuristicFlag = Accelerator.Allocate1D<int>(minCapacity);
+				}
+
 				BufferCapacity = minCapacity;
 			}
 
 			public void Dispose()
 			{
+				_stream?.Dispose();
+				_stream = null;
 				Return(this);
 			}
 
@@ -491,7 +449,7 @@ public sealed class PrimeTester
 				HeuristicFlag?.Dispose();
 				_stream?.Dispose();
 				_stream = null;
-				ReleaseSharedTables(_sharedTables);
+				// ReleaseSharedTables(_sharedTables);
 				// These resources are shared between GPU leases
 				// Accelerator.Dispose();
 				// AcceleratorContext.Dispose();
@@ -515,49 +473,49 @@ public sealed class PrimeTester
 
 		private sealed class SharedHeuristicGpuTables
 		{
+			private static readonly ConcurrentDictionary<Accelerator, SharedHeuristicGpuTables> _sharedTables = new(20_480, PerfectNumberConstants.RollingAccelerators);
+
 			private int _referenceCount;
 
+			internal static SharedHeuristicGpuTables EnsureStaticTables(Accelerator accelerator, AcceleratorStream stream)
+				=> _sharedTables.GetOrAdd(accelerator, _ => new SharedHeuristicGpuTables(accelerator, stream));
+
 			internal SharedHeuristicGpuTables(Accelerator accelerator, AcceleratorStream stream)
 			{
-
 				var primesDefault = DivisorGenerator.SmallPrimes;
-				DevicePrimesDefault = accelerator.Allocate1D<uint>(primesDefault.Length);
-				DevicePrimesDefault.View.CopyFromCPU(stream, primesDefault);
-
 				var primesLastOne = DivisorGenerator.SmallPrimesLastOne;
-				DevicePrimesLastOne = accelerator.Allocate1D<uint>(primesLastOne.Length);
-				DevicePrimesLastOne.View.CopyFromCPU(stream, primesLastOne);
-
 				var primesLastSeven = DivisorGenerator.SmallPrimesLastSeven;
-				DevicePrimesLastSeven = accelerator.Allocate1D<uint>(primesLastSeven.Length);
-				DevicePrimesLastSeven.View.CopyFromCPU(stream, primesLastSeven);
-
 				var primesLastThree = DivisorGenerator.SmallPrimesLastThree;
-				DevicePrimesLastThree = accelerator.Allocate1D<uint>(primesLastThree.Length);
-				DevicePrimesLastThree.View.CopyFromCPU(stream, primesLastThree);
-
 				var primesLastNine = DivisorGenerator.SmallPrimesLastNine;
-				DevicePrimesLastNine = accelerator.Allocate1D<uint>(primesLastNine.Length);
-				DevicePrimesLastNine.View.CopyFromCPU(stream, primesLastNine);
-
 				var primesPow2Default = DivisorGenerator.SmallPrimesPow2;
-				DevicePrimesPow2Default = accelerator.Allocate1D<ulong>(primesPow2Default.Length);
-				DevicePrimesPow2Default.View.CopyFromCPU(stream, primesPow2Default);
-
 				var primesPow2LastOne = DivisorGenerator.SmallPrimesPow2LastOne;
-				DevicePrimesPow2LastOne = accelerator.Allocate1D<ulong>(primesPow2LastOne.Length);
-				DevicePrimesPow2LastOne.View.CopyFromCPU(stream, primesPow2LastOne);
-
 				var primesPow2LastSeven = DivisorGenerator.SmallPrimesPow2LastSeven;
-				DevicePrimesPow2LastSeven = accelerator.Allocate1D<ulong>(primesPow2LastSeven.Length);
-				DevicePrimesPow2LastSeven.View.CopyFromCPU(stream, primesPow2LastSeven);
-
 				var primesPow2LastThree = DivisorGenerator.SmallPrimesPow2LastThree;
-				DevicePrimesPow2LastThree = accelerator.Allocate1D<ulong>(primesPow2LastThree.Length);
-				DevicePrimesPow2LastThree.View.CopyFromCPU(stream, primesPow2LastThree);
-
 				var primesPow2LastNine = DivisorGenerator.SmallPrimesPow2LastNine;
-				DevicePrimesPow2LastNine = accelerator.Allocate1D<ulong>(primesPow2LastNine.Length);
+
+				// lock (accelerator)
+				{
+					DevicePrimesDefault = accelerator.Allocate1D<uint>(primesDefault.Length);
+					DevicePrimesLastOne = accelerator.Allocate1D<uint>(primesLastOne.Length);
+					DevicePrimesLastSeven = accelerator.Allocate1D<uint>(primesLastSeven.Length);
+					DevicePrimesLastThree = accelerator.Allocate1D<uint>(primesLastThree.Length);
+					DevicePrimesLastNine = accelerator.Allocate1D<uint>(primesLastNine.Length);
+					DevicePrimesPow2Default = accelerator.Allocate1D<ulong>(primesPow2Default.Length);
+					DevicePrimesPow2LastOne = accelerator.Allocate1D<ulong>(primesPow2LastOne.Length);
+					DevicePrimesPow2LastSeven = accelerator.Allocate1D<ulong>(primesPow2LastSeven.Length);
+					DevicePrimesPow2LastThree = accelerator.Allocate1D<ulong>(primesPow2LastThree.Length);
+					DevicePrimesPow2LastNine = accelerator.Allocate1D<ulong>(primesPow2LastNine.Length);
+				}
+
+				DevicePrimesDefault.View.CopyFromCPU(stream, primesDefault);
+				DevicePrimesLastOne.View.CopyFromCPU(stream, primesLastOne);
+				DevicePrimesLastSeven.View.CopyFromCPU(stream, primesLastSeven);
+				DevicePrimesLastThree.View.CopyFromCPU(stream, primesLastThree);
+				DevicePrimesLastNine.View.CopyFromCPU(stream, primesLastNine);
+				DevicePrimesPow2Default.View.CopyFromCPU(stream, primesPow2Default);
+				DevicePrimesPow2LastOne.View.CopyFromCPU(stream, primesPow2LastOne);
+				DevicePrimesPow2LastSeven.View.CopyFromCPU(stream, primesPow2LastSeven);
+				DevicePrimesPow2LastThree.View.CopyFromCPU(stream, primesPow2LastThree);
 				DevicePrimesPow2LastNine.View.CopyFromCPU(stream, primesPow2LastNine);
 
 				HeuristicPrimeSieves.EnsureInitialized();
@@ -647,22 +605,20 @@ public sealed class PrimeTester
 					HeuristicGroupBDivisorSquaresEnding9.View);
 			}
 
-			internal void AddReference()
-			{
-				_referenceCount++;
-			}
+			[MethodImpl(MethodImplOptions.AggressiveInlining)]
+			internal void AddReference() => Atomic.Add(ref _referenceCount, 1);
 
-			internal void Release()
-			{
-				if (_referenceCount > 0)
-				{
-					_referenceCount--;
-				}
-			}
+			[MethodImpl(MethodImplOptions.AggressiveInlining)]
+			internal void Release() => Atomic.Add(ref _referenceCount, -1);
 
 			private static MemoryBuffer1D<ulong, Stride1D.Dense> CopySpanToDevice(Accelerator accelerator, AcceleratorStream stream, ReadOnlySpan<ulong> span)
 			{
-				var buffer = accelerator.Allocate1D<ulong>(span.Length);
+				MemoryBuffer1D<ulong, Stride1D.Dense>? buffer;
+				// lock (accelerator)
+				{
+					buffer = accelerator.Allocate1D<ulong>(span.Length);
+				}
+
 				if (!span.IsEmpty)
 				{
 					ref ulong sourceRef = ref MemoryMarshal.GetReference(span);
@@ -674,7 +630,12 @@ public sealed class PrimeTester
 
 			private static MemoryBuffer1D<ulong, Stride1D.Dense> CopyUintSpanToDevice(Accelerator accelerator, AcceleratorStream stream, ReadOnlySpan<uint> span)
 			{
-				var buffer = accelerator.Allocate1D<ulong>(span.Length);
+				MemoryBuffer1D<ulong, Stride1D.Dense>? buffer;
+				// lock (accelerator)
+				{
+					buffer = accelerator.Allocate1D<ulong>(span.Length);
+				}
+
 				if (!span.IsEmpty)
 				{
 					var converted = new ulong[span.Length];
@@ -722,9 +683,9 @@ public sealed class PrimeTester
 	}
 
 	// Expose cache clearing for accelerator disposal coordination
-	public static void ClearGpuCaches(Accelerator accelerator)
+	public static void ClearGpuCaches()
 	{
-		PrimeTesterGpuContextPool.Clear(accelerator);
+		PrimeTesterGpuContextPool.Clear();
 	}
 
 	internal static void DisposeGpuContexts()
@@ -746,6 +707,9 @@ public sealed class PrimeTester
 		return BinaryGcd(n, m) != 1UL;
 	}
 
+	[ThreadStatic]
+	private static readonly Dictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<byte>>>? _kernel;
+
 	internal static void SharesFactorWithMaxExponentBatch(ReadOnlySpan<ulong> values, Span<byte> results)
 	{
 		// TODO: Route this batch helper through the shared GPU kernel pool from
@@ -755,18 +719,24 @@ public sealed class PrimeTester
 		var accelerator = gpu.Accelerator;
 		var stream = gpu.Stream;
 
-		var kernel = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SharesFactorKernel)).CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<byte>>>();
-
 		int length = values.Length;
-		var inputBuffer = accelerator.Allocate1D<ulong>(length);
-		var resultBuffer = accelerator.Allocate1D<byte>(length);
-
 		ArrayPool<ulong> pool = ThreadStaticPools.UlongPool;
 		ulong[] temp = pool.Rent(length);
 		values.CopyTo(temp);
+
+		MemoryBuffer1D<ulong, Stride1D.Dense>? inputBuffer;
+		MemoryBuffer1D<byte, Stride1D.Dense>? resultBuffer;
+
+		// lock (accelerator)
+		{
+			inputBuffer = accelerator.Allocate1D<ulong>(length);
+			resultBuffer = accelerator.Allocate1D<byte>(length);
+		}
+
 		inputBuffer.View.CopyFromCPU(stream, ref temp[0], length);
+		var kernel = GetSharesFactorKernel(accelerator);
 		kernel(stream, length, inputBuffer.View, resultBuffer.View);
-		resultBuffer.View.CopyToCPU(stream, ref results[0], length);
+		resultBuffer.View.CopyToCPU(stream, in results);
 		stream.Synchronize();
 		pool.Return(temp, clearArray: false);
 		resultBuffer.Dispose();
@@ -774,6 +744,19 @@ public sealed class PrimeTester
 		gpu.Dispose();
 	}
 
+	private static Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<byte>> GetSharesFactorKernel(Accelerator accelerator)
+	{
+		var pool = _kernel ?? [];
+		if (pool.TryGetValue(accelerator, out var cached))
+		{
+			return cached;
+		}
+
+		cached = KernelUtil.GetKernel(accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SharesFactorKernel)).CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView<ulong>, ArrayView<byte>>>();
+		pool[accelerator] = cached;
+		return cached;
+	}
+
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	private static ulong BinaryGcd(ulong u, ulong v)
 	{

PerfectNumbers.Core/ULongExtensions.Gpu.cs

@@ -1,10 +1,10 @@
 using System;
 using System.Collections.Concurrent;
 using System.Runtime.CompilerServices;
-using ILGPU;
 using ILGPU.Algorithms;
 using ILGPU.Runtime;
 using PerfectNumbers.Core.Gpu;
+using PerfectNumbers.Core.Gpu.Accelerators;
 
 namespace PerfectNumbers.Core;
 
@@ -315,25 +315,16 @@ public static partial class ULongExtensions
 	}
 
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
-	public static ulong Pow2MontgomeryModWindowedGpu(this ulong exponent, in MontgomeryDivisorData divisor, bool keepMontgomery)
-	{
-		// We shouldn't hit it in production code. If we do, we're doing something wrong.
-		// ulong modulus = divisor.Modulus;
-		// if (exponent == 0UL)
-		// {
-		//      return divisor.MontgomeryOne;
-		// }
+	public static ulong Pow2MontgomeryModWindowedKeepGpu(this ulong exponent, in MontgomeryDivisorData divisor)
+		=> Pow2MontgomeryGpuCalculator.CalculateKeep(exponent, divisor);
 
-		// We barely ever hit it in production code. It's e.g. 772 calls out of billions
-		// if (exponent <= Pow2WindowFallbackThreshold)
-		// {
-		//      return Pow2MontgomeryModSingleBit(exponent, divisor);
-		// }
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static ulong Pow2MontgomeryModWindowedConvertGpu(this ulong exponent, in MontgomeryDivisorData divisor)
+		=> Pow2MontgomeryGpuCalculator.CalculateConvert(exponent, divisor);
 
-		return keepMontgomery
-				? Pow2MontgomeryGpuExecutor.ExecuteKeep(exponent, divisor)
-				: Pow2MontgomeryGpuExecutor.ExecuteConvert(exponent, divisor);
-	}
+	[MethodImpl(MethodImplOptions.AggressiveInlining)]
+	public static ulong Pow2MontgomeryModWindowedConvertGpu(this ulong exponent, Pow2MontgomeryAccelerator gpu, in MontgomeryDivisorData divisor)
+		=> Pow2MontgomeryGpuCalculator.CalculateConvert(gpu, exponent, divisor);
 
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	private static void InitializeStandardOddPowers(ulong modulus, Span<ulong> oddPowers)
@@ -387,16 +378,7 @@ public static partial class ULongExtensions
 	}
 
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
-	private static int GetPortableBitLengthGpu(ulong value)
-	{
-		// Keep this commented out. It will never happen in production code.
-		// if (value == 0UL)
-		// {
-		//         return 0;
-		// }
-
-		return 64 - XMath.LeadingZeroCount(value);
-	}
+	private static int GetPortableBitLengthGpu(ulong value) => 64 - XMath.LeadingZeroCount(value);
 
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	private static int GetWindowSizeGpu(int n)
@@ -429,102 +411,12 @@ public static partial class ULongExtensions
 	public static ulong Pow2MontgomeryModWithCycleGpu(this ulong exponent, ulong cycleLength, in MontgomeryDivisorData divisor)
 	{
 		ulong rotationCount = exponent % cycleLength;
-		return Pow2MontgomeryGpuExecutor.ExecuteConvert(rotationCount, divisor);
+		return Pow2MontgomeryGpuCalculator.CalculateConvert(rotationCount, divisor);
 	}
 
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	public static ulong Pow2MontgomeryModFromCycleRemainderGpu(this ulong reducedExponent, in MontgomeryDivisorData divisor)
 	{
-		return Pow2MontgomeryGpuExecutor.ExecuteConvert(reducedExponent, divisor);
-	}
-
-	private static class Pow2MontgomeryGpuExecutor
-	{
-		private static readonly Pow2MontgomeryKernelGroup Kernel = Pow2MontgomeryKernelGroup.Create(SharedGpuContext.Accelerator);
-
-		public static ulong ExecuteKeep(ulong exponent, in MontgomeryDivisorData divisor)
-		{
-			return Execute(exponent, divisor, useKeepKernel: true);
-		}
-
-		public static ulong ExecuteConvert(ulong exponent, in MontgomeryDivisorData divisor)
-		{
-			return Execute(exponent, divisor, useKeepKernel: false);
-		}
-
-		[ThreadStatic]
-		private static MemoryBuffer1D<ulong, Stride1D.Dense>? _exponentBuffer;
-		[ThreadStatic]
-		private static MemoryBuffer1D<ulong, Stride1D.Dense>? _resultBuffer;
-
-		private static ulong Execute(ulong exponent, in MontgomeryDivisorData divisor, bool useKeepKernel)
-		{
-			ulong result = 0UL;
-
-			GpuPrimeWorkLimiter.Acquire();
-
-			Accelerator accelerator = SharedGpuContext.Accelerator;
-			AcceleratorStream stream = accelerator.CreateStream();
-			var kernelGroup = Kernel;
-
-			EnsureCapacity();
-
-			MemoryBuffer1D<ulong, Stride1D.Dense> exponentBuffer = _exponentBuffer!;
-			MemoryBuffer1D<ulong, Stride1D.Dense> resultBuffer = _resultBuffer!;
-
-			exponentBuffer.View.CopyFromCPU(stream, ref exponent, 1);
-			// We don't need to worry about any left-overs here.
-			// resultBuffer.MemSetToZero(stream);
-
-			var kernel = useKeepKernel ? kernelGroup.KeepMontgomery : kernelGroup.ConvertToStandard;
-			kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View);
-			resultBuffer.View.CopyToCPU(stream, ref result, 1);
-			stream.Synchronize();
-			
-			stream.Dispose();
-			resultBuffer.Dispose();
-			exponentBuffer.Dispose();
-			GpuPrimeWorkLimiter.Release();
-			return result;
-		}
-
-		private static void EnsureCapacity()
-		{
-			if (_exponentBuffer == null)
-			{
-				Accelerator accelerator = SharedGpuContext.Accelerator;
-				_exponentBuffer = accelerator.Allocate1D<ulong>(1);
-				_resultBuffer = accelerator.Allocate1D<ulong>(1);
-			}
-		}
-
-		private sealed class Pow2MontgomeryKernelGroup
-		{
-			internal Pow2MontgomeryKernelGroup(
-					Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> keepMontgomery,
-					Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> convertToStandard)
-			{
-				KeepMontgomery = keepMontgomery;
-				ConvertToStandard = convertToStandard;
-			}
-
-			internal Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> KeepMontgomery { get; }
-
-			internal Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> ConvertToStandard { get; }
-
-			internal static Pow2MontgomeryKernelGroup Create(Accelerator accelerator)
-			{
-				var keepKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2MontgomeryKernels.Pow2MontgomeryKernelKeepMontgomery);
-				var keepLauncher = KernelUtil.GetKernel(keepKernel);
-				var keepDelegate = keepLauncher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
-
-				var convertKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2MontgomeryKernels.Pow2MontgomeryKernelConvertToStandard);
-				var convertLauncher = KernelUtil.GetKernel(convertKernel);
-				var convertDelegate = convertLauncher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
-
-				return new Pow2MontgomeryKernelGroup(keepDelegate, convertDelegate);
-			}
-		}
-
+		return Pow2MontgomeryGpuCalculator.CalculateConvert(reducedExponent, divisor);
 	}
 }