Various improvements to make sure scratch buffers don't override each others memory

PerfectNumbers.Core.Tests/Gpu/GpuPrimeLimiterTests.cs

@@ -18,19 +18,19 @@ public class GpuPrimeLimiterTests
         var startedSecond = new TaskCompletionSource<bool>();
         var sw = new Stopwatch();
 
-        var first = GpuPrimeWorkLimiter.Acquire();
+        GpuPrimeWorkLimiter.Acquire();
         sw.Start();
 
         var secondTask = Task.Run(() =>
         {
-            var second = GpuPrimeWorkLimiter.Acquire();
+            GpuPrimeWorkLimiter.Acquire();
             try
             {
                 startedSecond.TrySetResult(true);
             }
             finally
             {
-                second.Dispose();
+				GpuPrimeWorkLimiter.Release();
             }
         });
 
@@ -43,7 +43,7 @@ public class GpuPrimeLimiterTests
         finally
         {
             // Release first after the initial wait so the second task can proceed.
-            first.Dispose();
+			GpuPrimeWorkLimiter.Release();
         }
 
         await secondTask.WaitAsync(TimeSpan.FromSeconds(1));

PerfectNumbers.Core/Gpu/GpuContextPool.cs

@@ -14,11 +14,12 @@ public static class GpuContextPool
 		public readonly Context Context;
 		public readonly Accelerator Accelerator;
 
-		public PooledContext()
+		public PooledContext(KernelContainer kernels)
 		{
 			Context = SharedGpuContext.Context;
 			Accelerator = SharedGpuContext.Accelerator;
-			GpuStaticTableInitializer.EnsureStaticTables(Accelerator);
+			GpuStaticTableInitializer.EnsureStaticTables(kernels, Accelerator);
+			kernels.Dispose();
 			// NOTE: Avoid loading/compiling any kernel here to prevent implicit
 			// CL stream/queue creation during accelerator construction.
 			// Some OpenCL drivers are fragile when a queue is created immediately
@@ -52,7 +53,7 @@ public static class GpuContextPool
 		}
 
 		// Create a new accelerator when the pool does not have one available.
-		return new GpuContextLease(new PooledContext());
+		return new GpuContextLease(new PooledContext(GpuKernelPool.GetKernels()));
 	}
 
 
@@ -70,7 +71,8 @@ public static class GpuContextPool
 		PooledContext context;
 		for (int i = 0; i < toCreate; i++)
 		{
-			context = new PooledContext();
+			KernelContainer kernels = GpuKernelPool.GetKernels();
+			context = new PooledContext(kernels);
 			contexts[i] = context;
 			pool.Enqueue(context);
 		}

PerfectNumbers.Core/Gpu/GpuKernelLease.cs

@@ -11,25 +11,24 @@ public sealed class GpuKernelLease
 {
     private static readonly ConcurrentQueue<GpuKernelLease> Pool = new();
 
-    private Action? _releaseLimiter;
     private GpuContextLease? _gpu;
-    private KernelContainer? _kernels;
+    public KernelContainer Kernels;
     private AcceleratorStream? _stream;
 
-    private GpuKernelLease()
+    private GpuKernelLease(KernelContainer kernels)
 	{
+		Kernels = kernels;
 	}
 
-    internal static GpuKernelLease Rent(GpuPrimeWorkLimiter.Lease limiter, GpuContextLease gpu, KernelContainer kernels)
+    internal static GpuKernelLease Rent(GpuContextLease gpu, KernelContainer kernels)
     {
         if (!Pool.TryDequeue(out var lease))
         {
-            lease = new GpuKernelLease();
+            lease = new GpuKernelLease(kernels);
         }
 
-        lease._releaseLimiter = limiter.Dispose;
         lease._gpu = gpu;
-        lease._kernels = kernels;
+        lease.Kernels = kernels;
         lease._stream = null;
         return lease;
     }
@@ -57,7 +56,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator; // avoid capturing 'this' in lambda
-            return KernelContainer.InitOnce(ref _kernels!.Order, () =>
+            return KernelContainer.InitOnce(ref Kernels!.Order, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ulong, ArrayView<GpuUInt128>, ArrayView<ulong>>(OrderKernels.OrderKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -71,7 +70,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref _kernels!.Pow2Mod, () =>
+            return KernelContainer.InitOnce(ref Kernels!.Pow2Mod, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernels.Pow2ModKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -85,7 +84,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref _kernels!.Incremental, () =>
+            return KernelContainer.InitOnce(ref Kernels!.Incremental, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ulong, ulong, ulong, ulong, ArrayView<ulong>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -99,7 +98,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref _kernels!.IncrementalOrder, () =>
+            return KernelContainer.InitOnce(ref Kernels!.IncrementalOrder, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(IncrementalKernels.IncrementalOrderKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -113,7 +112,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref _kernels!.Pow2ModOrder, () =>
+            return KernelContainer.InitOnce(ref Kernels!.Pow2ModOrder, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong, ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>(Pow2ModKernels.Pow2ModOrderKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -122,16 +121,16 @@ public sealed class GpuKernelLease
         }
     }
 
-    public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> SpecialMaxKernel
+    public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>> SpecialMaxKernel
     {
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref _kernels!.SpecialMax, () =>
+            return KernelContainer.InitOnce(ref Kernels!.SpecialMax, () =>
             {
-                var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel);
+                var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel);
                 var kernel = KernelUtil.GetKernel(loaded);
-                return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>>();
+                return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>>>();
             });
         }
     }
@@ -141,7 +140,7 @@ public sealed class GpuKernelLease
         get
         {
             var accel = Accelerator;
-            return KernelContainer.InitOnce(ref _kernels!.SmallPrimeFactor, () =>
+            return KernelContainer.InitOnce(ref Kernels!.SmallPrimeFactor, () =>
             {
                 var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, uint, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>(SmallPrimeFactorKernels.SmallPrimeFactorKernelScan);
                 var kernel = KernelUtil.GetKernel(loaded);
@@ -158,10 +157,8 @@ public sealed class GpuKernelLease
         _gpu?.Dispose();
         _gpu = null;
 
-        _releaseLimiter?.Invoke();
+		Kernels?.Dispose();
-        _releaseLimiter = null;
+		GpuPrimeWorkLimiter.Release();
-
-        _kernels = null;
         Pool.Enqueue(this);
     }
 

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -38,7 +38,7 @@ public sealed class KernelContainer
 	public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
 		ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>? Pow2ModOrder;
 	public Action<AcceleratorStream, Index1D, ulong, uint, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>? SmallPrimeFactor;
-    public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>? SpecialMax;
+	public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ushort, Stride1D.Dense>>? SpecialMax;
 
 	// Optional device buffer with small divisor cycles (<= 4M). Index = divisor, value = cycle length.
 	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallCycles;
@@ -50,7 +50,7 @@ public sealed class KernelContainer
 	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimeFactorRemainingSlot;
 	public MemoryBuffer1D<ulong, Stride1D.Dense>? SpecialMaxFactors;
 	public MemoryBuffer1D<ulong, Stride1D.Dense>? SpecialMaxCandidates;
-    public MemoryBuffer1D<byte, Stride1D.Dense>? SpecialMaxResult;
+	public MemoryBuffer1D<ushort, Stride1D.Dense>? SpecialMaxResult;
 	public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastOne;
 	public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimesPow2LastOne;
 	public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastSeven;
@@ -72,6 +72,8 @@ public sealed class KernelContainer
 		}
 		return current;
 	}
+
+	public void Dispose() => GpuKernelPool.Return(this);
 }
 
 public readonly struct ResiduePrimeViews(
@@ -86,44 +88,29 @@ public readonly struct ResiduePrimeViews(
 	public readonly ArrayView1D<ulong, Stride1D.Dense> LastSevenPow2 = lastSevenPow2;
 }
 
-public readonly struct SmallPrimeFactorTables
+public readonly struct SmallPrimeFactorTables(
-{
-    public readonly MemoryBuffer1D<uint, Stride1D.Dense> Primes;
-    public readonly MemoryBuffer1D<ulong, Stride1D.Dense> Squares;
-    public readonly int Count;
-
-    public SmallPrimeFactorTables(
 	MemoryBuffer1D<uint, Stride1D.Dense> primes,
 	MemoryBuffer1D<ulong, Stride1D.Dense> squares,
 	int count)
-    {
+{
-        Primes = primes;
+	public readonly MemoryBuffer1D<uint, Stride1D.Dense> Primes = primes;
-        Squares = squares;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> Squares = squares;
-        Count = count;
+	public readonly int Count = count;
-    }
 
 	public ArrayView1D<uint, Stride1D.Dense> PrimesView => Primes.View;
 	public ArrayView1D<ulong, Stride1D.Dense> SquaresView => Squares.View;
 }
 
-public readonly struct SmallPrimeFactorScratch
+public readonly struct SmallPrimeFactorScratch(
-{
-    public readonly MemoryBuffer1D<ulong, Stride1D.Dense> PrimeSlots;
-    public readonly MemoryBuffer1D<int, Stride1D.Dense> ExponentSlots;
-    public readonly MemoryBuffer1D<int, Stride1D.Dense> CountSlot;
-    public readonly MemoryBuffer1D<ulong, Stride1D.Dense> RemainingSlot;
-
-    public SmallPrimeFactorScratch(
 		MemoryBuffer1D<ulong, Stride1D.Dense> primeSlots,
 		MemoryBuffer1D<int, Stride1D.Dense> exponentSlots,
 		MemoryBuffer1D<int, Stride1D.Dense> countSlot,
 		MemoryBuffer1D<ulong, Stride1D.Dense> remainingSlot)
-    {
+{
-        PrimeSlots = primeSlots;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> PrimeSlots = primeSlots;
-        ExponentSlots = exponentSlots;
+	public readonly MemoryBuffer1D<int, Stride1D.Dense> ExponentSlots = exponentSlots;
-        CountSlot = countSlot;
+	public readonly MemoryBuffer1D<int, Stride1D.Dense> CountSlot = countSlot;
-        RemainingSlot = remainingSlot;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> RemainingSlot = remainingSlot;
-    }
 
 	public void Clear()
 	{
@@ -134,74 +121,55 @@ public readonly struct SmallPrimeFactorScratch
 	}
 }
 
-public readonly struct SpecialMaxScratch
+public readonly struct SpecialMaxScratch(
-{
-    public readonly MemoryBuffer1D<ulong, Stride1D.Dense> FactorValues;
-    public readonly MemoryBuffer1D<ulong, Stride1D.Dense> CandidateValues;
-    public readonly MemoryBuffer1D<byte, Stride1D.Dense> ResultSlot;
-
-    public SpecialMaxScratch(
 	MemoryBuffer1D<ulong, Stride1D.Dense> factorValues,
 	MemoryBuffer1D<ulong, Stride1D.Dense> candidateValues,
-        MemoryBuffer1D<byte, Stride1D.Dense> resultSlot)
+	MemoryBuffer1D<ushort, Stride1D.Dense> resultSlot)
-    {
+{
-        FactorValues = factorValues;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> FactorValues = factorValues;
-        CandidateValues = candidateValues;
+	public readonly MemoryBuffer1D<ulong, Stride1D.Dense> CandidateValues = candidateValues;
-        ResultSlot = resultSlot;
+	public readonly MemoryBuffer1D<ushort, Stride1D.Dense> ResultSlot = resultSlot;
-    }
 
 	public ArrayView1D<ulong, Stride1D.Dense> FactorsView => FactorValues.View;
-
 	public ArrayView1D<ulong, Stride1D.Dense> CandidatesView => CandidateValues.View;
-
+	public ArrayView1D<ushort, Stride1D.Dense> ResultView => ResultSlot.View;
-    public ArrayView1D<byte, Stride1D.Dense> ResultView => ResultSlot.View;
 }
 
 public class GpuKernelPool
 {
-    private static readonly ConcurrentDictionary<Accelerator, KernelContainer> KernelCache = new(); // TODO: Replace this concurrent map with a simple accelerator-indexed lookup once kernel launchers are prewarmed during startup so we can drop the thread-safe wrapper entirely.
+	private static readonly ConcurrentQueue<KernelContainer> KernelCache = new(); // TODO: Replace this concurrent map with a simple accelerator-indexed lookup once kernel launchers are prewarmed during startup so we can drop the thread-safe wrapper entirely.
 
-    private static KernelContainer GetKernels(Accelerator accelerator)
+	public static KernelContainer GetKernels() => KernelCache.TryDequeue(out var container)
-    {
+		? container
-        return KernelCache.GetOrAdd(accelerator, _ => new KernelContainer());
+		: new KernelContainer();
-    }
 
-    internal static void PreloadStaticTables(Accelerator accelerator)
+	public static void Return(KernelContainer container) => KernelCache.Enqueue(container);
+
+	internal static void PreloadStaticTables(KernelContainer kernels, Accelerator accelerator)
 	{
-        EnsureSmallCyclesOnDevice(accelerator);
+		EnsureSmallCyclesOnDevice(kernels, accelerator);
-        _ = EnsureSmallPrimesOnDevice(accelerator);
+		EnsureSmallPrimesOnDevice(kernels, accelerator);
-        _ = EnsureSmallPrimeFactorTables(accelerator);
+		EnsureSmallPrimeFactorTables(kernels, accelerator);
 	}
 
 	// 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).
-    public static ArrayView1D<ulong, Stride1D.Dense> EnsureSmallCyclesOnDevice(Accelerator accelerator)
+	public static ArrayView1D<ulong, Stride1D.Dense> EnsureSmallCyclesOnDevice(KernelContainer kernels, Accelerator accelerator)
 	{
-        var kernels = GetKernels(accelerator);
 		if (kernels.SmallCycles is { } buffer)
 		{
 			return buffer.View;
 		}
 
-        // Ensure single upload per accelerator even if multiple threads race here.
-        lock (kernels) // TODO: Remove this lock by pre-uploading the immutable small-cycle snapshot during initialization; once no mutation happens at runtime, the pool must expose a simple reference without synchronization.
-        {
-            if (kernels.SmallCycles is { } existing)
-            {
-                return existing.View;
-            }
-
 		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 = accelerator.Allocate1D<ulong>(host.Length);
 		device.View.CopyFromCPU(host);
 		kernels.SmallCycles = device;
 		return device.View;
 	}
-    }
 
-    public static ResiduePrimeViews EnsureSmallPrimesOnDevice(Accelerator accelerator)
+	public static ResiduePrimeViews EnsureSmallPrimesOnDevice(KernelContainer kernels, Accelerator accelerator)
 	{
-        var kernels = GetKernels(accelerator);
 		if (kernels.SmallPrimesLastOne is { } lastOne &&
 			kernels.SmallPrimesLastSeven is { } lastSeven &&
 			kernels.SmallPrimesPow2LastOne is { } lastOnePow2 &&
@@ -210,8 +178,6 @@ public class GpuKernelPool
 			return new ResiduePrimeViews(lastOne.View, lastSeven.View, lastOnePow2.View, lastSevenPow2.View);
 		}
 
-        lock (kernels) // TODO: Inline these small-prime uploads into startup initialization alongside the small-cycle snapshot so we can drop runtime locking and keep the GPU pool free of synchronization.
-        {
 		if (kernels.SmallPrimesLastOne is { } existingLastOne &&
 			kernels.SmallPrimesLastSeven is { } existingLastSeven &&
 			kernels.SmallPrimesPow2LastOne is { } existingLastOnePow2 &&
@@ -241,18 +207,14 @@ public class GpuKernelPool
 
 		return new ResiduePrimeViews(deviceLastOne.View, deviceLastSeven.View, deviceLastOnePow2.View, deviceLastSevenPow2.View);
 	}
-    }
 
-    public static SmallPrimeFactorTables EnsureSmallPrimeFactorTables(Accelerator accelerator)
+	public static SmallPrimeFactorTables EnsureSmallPrimeFactorTables(KernelContainer kernels, Accelerator accelerator)
 	{
-        var kernels = GetKernels(accelerator);
 		if (kernels.SmallPrimeFactorsPrimes is { } primeBuffer && kernels.SmallPrimeFactorsSquares is { } squareBuffer)
 		{
 			return new SmallPrimeFactorTables(primeBuffer, squareBuffer, (int)primeBuffer.Length);
 		}
 
-        lock (kernels)
-        {
 		if (kernels.SmallPrimeFactorsPrimes is { } existingPrimes && kernels.SmallPrimeFactorsSquares is { } existingSquares)
 		{
 			return new SmallPrimeFactorTables(existingPrimes, existingSquares, (int)existingPrimes.Length);
@@ -271,12 +233,8 @@ public class GpuKernelPool
 
 		return new SmallPrimeFactorTables(devicePrimes, deviceSquares, hostPrimes.Length);
 	}
-    }
 
-    public static SmallPrimeFactorScratch EnsureSmallPrimeFactorScratch(Accelerator accelerator, int slotCount)
+	public static SmallPrimeFactorScratch EnsureSmallPrimeFactorScratch(KernelContainer kernels, Accelerator accelerator, int slotCount)
-    {
-        var kernels = GetKernels(accelerator);
-        lock (kernels)
 	{
 		MemoryBuffer1D<ulong, Stride1D.Dense>? primeSlots = kernels.SmallPrimeFactorPrimeSlots;
 		if (primeSlots is null || primeSlots.Length < slotCount)
@@ -310,17 +268,8 @@ public class GpuKernelPool
 
 		return new SmallPrimeFactorScratch(primeSlots!, exponentSlots!, countSlot!, remainingSlot!);
 	}
-    }
 
-    public static SpecialMaxScratch EnsureSpecialMaxScratch(Accelerator accelerator, int factorCapacity)
+	public static SpecialMaxScratch EnsureSpecialMaxScratch(KernelContainer kernels, Accelerator accelerator, int factorCapacity)
-    {
-        if (factorCapacity <= 0)
-        {
-            factorCapacity = 1;
-        }
-
-        var kernels = GetKernels(accelerator);
-        lock (kernels)
 	{
 		MemoryBuffer1D<ulong, Stride1D.Dense>? factorValues = kernels.SpecialMaxFactors;
 		if (factorValues is null || factorValues.Length < factorCapacity)
@@ -338,24 +287,22 @@ public class GpuKernelPool
 			kernels.SpecialMaxCandidates = candidateValues;
 		}
 
-            MemoryBuffer1D<byte, Stride1D.Dense>? resultSlot = kernels.SpecialMaxResult;
+		MemoryBuffer1D<ushort, Stride1D.Dense>? resultSlot = kernels.SpecialMaxResult;
 		if (resultSlot is null)
 		{
-                resultSlot = accelerator.Allocate1D<byte>(1);
+			resultSlot = accelerator.Allocate1D<ushort>(1);
 			kernels.SpecialMaxResult = resultSlot;
 		}
 
 		return new SpecialMaxScratch(factorValues!, candidateValues!, resultSlot!);
 	}
-    }
 
 	public static GpuKernelLease GetKernel(bool useGpuOrder)
 	{
-        var limiter = GpuPrimeWorkLimiter.Acquire();
+		GpuPrimeWorkLimiter.Acquire();
 		var gpu = Rent();
-        var accelerator = gpu.Accelerator;
+		var kernels = GetKernels();
-        var kernels = GetKernels(accelerator);
+		return GpuKernelLease.Rent(gpu, kernels);
-        return GpuKernelLease.Rent(limiter, gpu, kernels);
 	}
 
 	/// <summary>

PerfectNumbers.Core/Gpu/GpuPrimeWorkLimiter.cs

@@ -8,14 +8,8 @@ public static class GpuPrimeWorkLimiter
     private static SemaphoreSlim _semaphore = new(1, int.MaxValue);
     private static int _currentLimit = 1;
 
-    public static Lease Acquire()
+	public static void Acquire() => _semaphore.Wait();
-    {
+	public static void Release() => _semaphore.Release();
-        // TODO: Inline the pooled limiter guard from the GpuLimiterThroughputBenchmarks so prime-sieve
-        // GPU jobs stop allocating Lease instances on every acquisition.
-        var sem = _semaphore;
-        sem.Wait();
-        return new Lease(sem);
-    }
 
 	public static void SetLimit(int value)
     {
@@ -36,19 +30,4 @@ public static class GpuPrimeWorkLimiter
         _currentLimit = value;
         previous.Dispose();
     }
-
-    public sealed class Lease : IDisposable
-    {
-        private readonly SemaphoreSlim _sem;
-
-        internal Lease(SemaphoreSlim sem)
-        {
-            _sem = sem;
-        }
-
-        public void Dispose()
-        {
-            _sem.Release();
-        }
-    }
 }

PerfectNumbers.Core/Gpu/GpuStaticTableInitializer.cs

@@ -11,18 +11,13 @@ internal static class GpuStaticTableInitializer
 {
     private static readonly ConcurrentDictionary<Accelerator, byte> Initialized = new(AcceleratorReferenceComparer.Instance);
 
-    internal static void EnsureStaticTables(Accelerator accelerator)
+    internal static void EnsureStaticTables(KernelContainer kernels, Accelerator accelerator)
     {
-        if (accelerator is null)
+        Initialized.GetOrAdd(accelerator, acc =>
         {
-            return;
+            PrimeTester.PrimeTesterGpuContextPool.EnsureStaticTables(kernels, acc);
-        }
+            GpuKernelPool.PreloadStaticTables(kernels, acc);
-
+            PrimeOrderGpuHeuristics.PreloadStaticTables( acc);
-        Initialized.GetOrAdd(accelerator, static acc =>
-        {
-            PrimeTester.PrimeTesterGpuContextPool.EnsureStaticTables(acc);
-            GpuKernelPool.PreloadStaticTables(acc);
-            PrimeOrderGpuHeuristics.PreloadStaticTables(acc);
             return 0;
         });
     }

PerfectNumbers.Core/Gpu/Kernels/PrimeOrderSpecialMaxKernels.cs

@@ -13,7 +13,7 @@ internal static partial class PrimeOrderGpuHeuristics
         int factorCount,
         MontgomeryDivisorData divisor,
         ArrayView1D<ulong, Stride1D.Dense> candidates,
-        ArrayView1D<byte, Stride1D.Dense> resultOut)
+        ArrayView1D<ushort, Stride1D.Dense> resultOut)
     {
         if (index != 0)
         {

PerfectNumbers.Core/Gpu/MersenneNumberIncrementalGpuTester.cs

@@ -31,11 +31,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.EnsureSmallCyclesOnDevice(accelerator);
+        var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(gpuLease.Kernels, accelerator);
         ResiduePrimeViews primeViews = default;
         if (_kernelType == GpuKernelType.Pow2Mod)
         {
-            primeViews = GpuKernelPool.EnsureSmallPrimesOnDevice(accelerator);
+            primeViews = GpuKernelPool.EnsureSmallPrimesOnDevice(gpuLease.Kernels, accelerator);
         }
 
         var orderBuffer = accelerator.Allocate1D<ulong>(batchSize);

PerfectNumbers.Core/Gpu/MersenneNumberLucasLehmerGpuTester.cs

@@ -54,13 +54,13 @@ public class MersenneNumberLucasLehmerGpuTester
             return false;
         }
 
-        var limiter = GpuPrimeWorkLimiter.Acquire();
+        GpuPrimeWorkLimiter.Acquire();
         bool result;
         if (exponent >= 128UL)
         {
             if (!TryGetNttParameters(exponent, out var nttMod, out var primitiveRoot))
             {
-                limiter.Dispose();
+				GpuPrimeWorkLimiter.Release();
                 throw new NotSupportedException("NTT parameters for the given exponent are not available.");
             }
 
@@ -80,7 +80,7 @@ public class MersenneNumberLucasLehmerGpuTester
             gpu.Dispose();
         }
 
-        limiter.Dispose();
+		GpuPrimeWorkLimiter.Release();
         return result;
     }
 

PerfectNumbers.Core/Gpu/MersenneNumberOrderGpuTester.cs

@@ -29,6 +29,8 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
 		UInt128 remaining;
 		int currentSize;
 		int found = 0;
+
+		var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(gpuLease.Kernels, accelerator);
 		while (kStart <= maxK && Volatile.Read(ref isPrime))
 		{
 			remaining = maxK - kStart + 1UL;
@@ -44,7 +46,6 @@ public class MersenneNumberOrderGpuTester(GpuKernelType kernelType, bool useGpuO
 			var kernelArgs = new ResidueAutomatonArgs(q0m10, step10, q0m8, step8, q0m3, step3, q0m5, step5);
 
 			// Ensure device has small cycles table for early rejections in order kernels
-                        var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(accelerator);
 			kernel(stream, currentSize, exponent, (GpuUInt128)twoP, (GpuUInt128)kStart, last, divMul,
 							kernelArgs, foundBuffer.View, smallCyclesView);
 

PerfectNumbers.Core/Gpu/MersenneNumberResidueGpuTester.cs

@@ -26,7 +26,7 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
     // 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)
     {
-        var limiter = GpuPrimeWorkLimiter.Acquire();
+        GpuPrimeWorkLimiter.Acquire();
         var gpuLease = RentAccelerator();
         var accelerator = gpuLease.Accelerator;
 
@@ -35,14 +35,15 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
         var stream = accelerator.CreateStream();
         var kernel = GetPow2ModKernel(accelerator);
 		var resources = RentResources(accelerator, GpuConstants.ScanBatchSize);
+		var gpuKernels = GpuKernelPool.GetKernels();
 
         var orderBuffer = resources.OrderBuffer;
         ulong[] orderArray = resources.OrderArray;
         int batchSize = resources.Capacity;
 
         // Ensure device has small cycles and primes tables for in-kernel lookup
-        var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(accelerator);
+        var smallCyclesView = GpuKernelPool.EnsureSmallCyclesOnDevice(gpuKernels, accelerator);
-        ResiduePrimeViews primeViews = GpuKernelPool.EnsureSmallPrimesOnDevice(accelerator);
+        ResiduePrimeViews primeViews = GpuKernelPool.EnsureSmallPrimesOnDevice(gpuKernels, accelerator);
 
         UInt128 kStart = 1UL;
         UInt128 limit = maxK + UInt128.One;
@@ -109,8 +110,9 @@ public class MersenneNumberResidueGpuTester(bool useGpuOrder)
         ReturnResources(accelerator, resources);
         stream.Dispose();
 		// Monitor.Exit(gpuLease.ExecutionLock);
+		gpuKernels.Dispose();
 		ReturnAccelerator(gpuLease);
-        limiter.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) =>

PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs

@@ -1,5 +1,4 @@
 using Open.Numeric.Primes;
-using System;
 using ILGPU.Runtime;
 using System.Runtime.CompilerServices;
 using PerfectNumbers.Core.Gpu;
@@ -225,10 +224,10 @@ public sealed class HeuristicCombinedPrimeTester
 
 	private static bool HeuristicTrialDivisionGpuDetectsDivisor(ulong n, ulong maxDivisorSquare, byte nMod10)
 	{
-                var limiter = GpuPrimeWorkLimiter.Acquire();
+		GpuPrimeWorkLimiter.Acquire();
 		var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(1);
 		var accelerator = gpu.Accelerator;
-                var kernel = gpu.HeuristicTrialDivisionKernel;
+		var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.HeuristicTrialDivisionKernel;
 		var flagView1D = gpu.HeuristicFlag.View;
 		var flagView = flagView1D.AsContiguous();
 
@@ -249,7 +248,7 @@ public sealed class HeuristicCombinedPrimeTester
 		compositeDetected = compositeFlag != 0;
 
 		gpu.Dispose();
-                limiter.Dispose();
+		GpuPrimeWorkLimiter.Release();
 		return compositeDetected;
 	}
 

PerfectNumbers.Core/HeuristicPrimeTester.cs

@@ -240,10 +240,10 @@ public sealed class HeuristicPrimeTester
 
 	private bool HeuristicTrialDivisionGpuDetectsDivisor(ulong n, ulong maxDivisorSquare, byte nMod10)
 	{
-                var limiter = GpuPrimeWorkLimiter.Acquire();
+		GpuPrimeWorkLimiter.Acquire();
 		var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(1);
 		var accelerator = gpu.Accelerator;
-                var kernel = gpu.HeuristicTrialDivisionKernel;
+		var kernel = PrimeTester.PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.HeuristicTrialDivisionKernel;
 		var flagView1D = gpu.HeuristicFlag.View;
 		var flagView = flagView1D.AsContiguous();
 
@@ -301,7 +301,7 @@ public sealed class HeuristicPrimeTester
 		}
 
 		gpu.Dispose();
-                limiter.Dispose();
+		GpuPrimeWorkLimiter.Release();
 		return compositeDetected;
 	}
 

PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs

@@ -3,6 +3,7 @@ using System.Diagnostics;
 using System.Numerics;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+using ILGPU;
 using PerfectNumbers.Core.Gpu;
 
 namespace PerfectNumbers.Core;
@@ -106,6 +107,7 @@ internal static partial class PrimeOrderCalculator
 		return candidateOrder;
 	}
 
+	// private static ulong specialMaxHits;
 	private static bool TrySpecialMaxCpu(ulong phi, ulong prime, PartialFactorResult factors, in MontgomeryDivisorData divisorData)
 	{
 		int length = factors.Count;
@@ -125,6 +127,8 @@ internal static partial class PrimeOrderCalculator
 			return EvaluateSpecialMaxCandidates(stackBuffer, factorSpan, phi, prime, divisorData);
 		}
 
+		// Atomic.Add(ref specialMaxHits, 1);
+		// Console.WriteLine($"Special max GPU hit {specialMaxHits} ({length})");
 		return EvaluateSpecialMaxCandidatesGpu(stackBuffer, factorSpan, phi, prime, divisorData);
 	}
 

PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs

@@ -25,8 +25,8 @@ internal static partial class PrimeOrderCalculator
 		Accelerator accelerator = lease.Accelerator;
 		AcceleratorStream stream = lease.Stream;
 
-		SmallPrimeFactorTables tables = GpuKernelPool.EnsureSmallPrimeFactorTables(accelerator);
+		SmallPrimeFactorTables tables = GpuKernelPool.EnsureSmallPrimeFactorTables(lease.Kernels, accelerator);
-		SmallPrimeFactorScratch scratch = GpuKernelPool.EnsureSmallPrimeFactorScratch(accelerator, GpuSmallPrimeFactorSlots);
+		SmallPrimeFactorScratch scratch = GpuKernelPool.EnsureSmallPrimeFactorScratch(lease.Kernels, accelerator, GpuSmallPrimeFactorSlots);
 
 		var kernel = lease.SmallPrimeFactorKernel;
 		kernel(
@@ -92,7 +92,7 @@ internal static partial class PrimeOrderCalculator
 		Accelerator accelerator = lease.Accelerator;
 		AcceleratorStream stream = lease.Stream;
 
-		SpecialMaxScratch scratch = GpuKernelPool.EnsureSpecialMaxScratch(accelerator, factorCount);
+		SpecialMaxScratch scratch = GpuKernelPool.EnsureSpecialMaxScratch(lease.Kernels, accelerator, factorCapacity: factorCount);
 
 		scratch.FactorsView.SubView(0, factorCount).CopyFromCPU(ref MemoryMarshal.GetReference(factorSpan), factorCount);
 
@@ -109,9 +109,8 @@ internal static partial class PrimeOrderCalculator
 
 		stream.Synchronize();
 
-		byte result = 0;
+		ushort result = 0;
 		scratch.ResultView.CopyToCPU(ref result, 1);
-
 		lease.Dispose();
 
 		return result != 0;

PerfectNumbers.Core/PrimeTester.cs

@@ -95,10 +95,10 @@ public sealed class PrimeTester
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	public static bool IsPrimeGpu(ulong n)
 	{
-		var limiter = GpuPrimeWorkLimiter.Acquire();
+		GpuPrimeWorkLimiter.Acquire();
 		var gpu = PrimeTesterGpuContextPool.Rent(1);
 		var accelerator = gpu.Accelerator;
-		var kernel = gpu.Kernel;
+		var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.Kernel;
 
 		ulong value = n;
 		byte flag = 0;
@@ -127,7 +127,7 @@ public sealed class PrimeTester
 		outputView.CopyToCPU(ref flag, 1);
 
 		gpu.Dispose();
-		limiter.Dispose();
+		GpuPrimeWorkLimiter.Release();
 
 		return flag != 0;
 	}
@@ -160,10 +160,10 @@ public sealed class PrimeTester
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	public static void IsPrimeBatchGpu(ReadOnlySpan<ulong> values, Span<byte> results)
 	{
-		var limiter = GpuPrimeWorkLimiter.Acquire();
+		GpuPrimeWorkLimiter.Acquire();
 		var gpu = PrimeTesterGpuContextPool.Rent(GpuBatchSize);
 		var accelerator = gpu.Accelerator;
-		var kernel = gpu.Kernel;
+		var kernel = PrimeTesterGpuContextPool.PrimeTesterGpuContextLease.Kernel;
 		int totalLength = values.Length;
 		int batchSize = GpuBatchSize;
 
@@ -205,7 +205,7 @@ public sealed class PrimeTester
 		}
 
 		gpu.Dispose();
-		limiter.Dispose();
+		GpuPrimeWorkLimiter.Release();
 	}
 
 	internal static class PrimeTesterGpuContextPool
@@ -244,13 +244,8 @@ public sealed class PrimeTester
 			}
 		}
 
-		internal static void EnsureStaticTables(Accelerator accelerator)
+		internal static void EnsureStaticTables(KernelContainer kernels, Accelerator accelerator)
 		{
-			if (accelerator is null)
-			{
-				return;
-			}
-
 			var tables = RentSharedTables(accelerator);
 			var heuristicTables = tables.CreateHeuristicDivisorTables();
 			HeuristicGpuDivisorTables.InitializeShared(in heuristicTables);
@@ -271,8 +266,8 @@ public sealed class PrimeTester
 					return;
 				}
 
-				Accelerator accelerator = SharedGpuContext.Accelerator;
+				KernelContainer kernels = GpuKernelPool.GetKernels();
-				GpuStaticTableInitializer.EnsureStaticTables(accelerator);
+				GpuStaticTableInitializer.EnsureStaticTables(kernels, SharedGpuContext.Accelerator);
 
 				int toCreate = target - WarmedLeaseCount;
 				for (int i = 0; i < toCreate; i++)
@@ -280,6 +275,7 @@ public sealed class PrimeTester
 					Pool.Enqueue(new PrimeTesterGpuContextLease(minBufferCapacity));
 				}
 
+				kernels.Dispose();
 				WarmedLeaseCount = target;
 			}
 		}
@@ -383,8 +379,11 @@ public sealed class PrimeTester
 
 			internal readonly Context AcceleratorContext;
 			public readonly Accelerator Accelerator;
-			public readonly Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> Kernel;
+			
-			public readonly Action<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables> HeuristicTrialDivisionKernel;
+			public static readonly Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> Kernel = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
+
+			public static readonly Action<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables> HeuristicTrialDivisionKernel = SharedGpuContext.Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel);
+			
 			public MemoryBuffer1D<ulong, Stride1D.Dense> Input = null!;
 			public MemoryBuffer1D<byte, Stride1D.Dense> Output = null!;
 			public MemoryBuffer1D<int, Stride1D.Dense> HeuristicFlag = null!;
@@ -423,9 +422,9 @@ public sealed class PrimeTester
 			{
 				AcceleratorContext = SharedGpuContext.Context;
 				Accelerator = SharedGpuContext.Accelerator;
-				GpuStaticTableInitializer.EnsureStaticTables(Accelerator);
+				// GpuStaticTableInitializer.EnsureStaticTables(Accelerator);
-				Kernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
+				// Kernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<uint>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
-				HeuristicTrialDivisionKernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel);
+				// HeuristicTrialDivisionKernel = Accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<int>, ulong, ulong, HeuristicGpuDivisorTableKind, HeuristicGpuDivisorTables>(PrimeTesterKernels.HeuristicTrialDivisionKernel);
 
 				_sharedTables = RentSharedTables(Accelerator);