GPU stabilization improvements to prevent calculation crashes

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);
 }