Use divisor squares to bound heuristic combined trial division

PerfectNumbers.Core/Gpu/Kernels/PrimeTesterKernels.cs

@@ -68,16 +68,18 @@ internal static class PrimeTesterKernels
         results[index] = result;
     }
 
-    public static void HeuristicTrialDivisionKernel(Index1D index, ArrayView<ulong> divisors, ulong n, ArrayView<byte> results)
+    public static void HeuristicTrialDivisionKernel(Index1D index, ArrayView<ulong> divisors, ulong n, ArrayView<byte> resultFlag)
     {
         ulong divisor = divisors[index];
         if (divisor <= 1UL)
         {
-            results[index] = 0;
             return;
         }
 
-        results[index] = n % divisor == 0UL ? (byte)1 : (byte)0;
+        if (n % divisor == 0UL)
+        {
+            resultFlag[0] = 1;
+        }
     }
 
     public static void SharesFactorKernel(Index1D index, ArrayView<ulong> numbers, ArrayView<byte> results)

PerfectNumbers.Core/HeuristicCombinedPrimeTester.cs

@@ -17,29 +17,33 @@ public sealed class HeuristicCombinedPrimeTester
 	{
 		HeuristicPrimeSieves.EnsureInitialized();
 
-		CombinedDivisorsEnding1TwoAOneB = BuildCombinedDivisors(1, CombinedDivisorPattern.TwoAOneB);
+		CombinedDivisorsEnding1TwoAOneB = BuildCombinedDivisors(1, CombinedDivisorPattern.TwoAOneB, out CombinedDivisorsEnding1TwoAOneBSquares);
-		CombinedDivisorsEnding3TwoAOneB = BuildCombinedDivisors(3, CombinedDivisorPattern.TwoAOneB);
+		CombinedDivisorsEnding3TwoAOneB = BuildCombinedDivisors(3, CombinedDivisorPattern.TwoAOneB, out CombinedDivisorsEnding3TwoAOneBSquares);
-		CombinedDivisorsEnding7TwoAOneB = BuildCombinedDivisors(7, CombinedDivisorPattern.TwoAOneB);
+		CombinedDivisorsEnding7TwoAOneB = BuildCombinedDivisors(7, CombinedDivisorPattern.TwoAOneB, out CombinedDivisorsEnding7TwoAOneBSquares);
-		CombinedDivisorsEnding9TwoAOneB = BuildCombinedDivisors(9, CombinedDivisorPattern.TwoAOneB);
+		CombinedDivisorsEnding9TwoAOneB = BuildCombinedDivisors(9, CombinedDivisorPattern.TwoAOneB, out CombinedDivisorsEnding9TwoAOneBSquares);
 
-		CombinedDivisorsEnding1OneAOneB = BuildCombinedDivisors(1, CombinedDivisorPattern.OneAOneB);
+		CombinedDivisorsEnding1OneAOneB = BuildCombinedDivisors(1, CombinedDivisorPattern.OneAOneB, out CombinedDivisorsEnding1OneAOneBSquares);
-		CombinedDivisorsEnding3OneAOneB = BuildCombinedDivisors(3, CombinedDivisorPattern.OneAOneB);
+		CombinedDivisorsEnding3OneAOneB = BuildCombinedDivisors(3, CombinedDivisorPattern.OneAOneB, out CombinedDivisorsEnding3OneAOneBSquares);
-		CombinedDivisorsEnding7OneAOneB = BuildCombinedDivisors(7, CombinedDivisorPattern.OneAOneB);
+		CombinedDivisorsEnding7OneAOneB = BuildCombinedDivisors(7, CombinedDivisorPattern.OneAOneB, out CombinedDivisorsEnding7OneAOneBSquares);
-		CombinedDivisorsEnding9OneAOneB = BuildCombinedDivisors(9, CombinedDivisorPattern.OneAOneB);
+		CombinedDivisorsEnding9OneAOneB = BuildCombinedDivisors(9, CombinedDivisorPattern.OneAOneB, out CombinedDivisorsEnding9OneAOneBSquares);
 
-		CombinedDivisorsEnding1ThreeAOneB = BuildCombinedDivisors(1, CombinedDivisorPattern.ThreeAOneB);
+		CombinedDivisorsEnding1ThreeAOneB = BuildCombinedDivisors(1, CombinedDivisorPattern.ThreeAOneB, out CombinedDivisorsEnding1ThreeAOneBSquares);
-		CombinedDivisorsEnding3ThreeAOneB = BuildCombinedDivisors(3, CombinedDivisorPattern.ThreeAOneB);
+		CombinedDivisorsEnding3ThreeAOneB = BuildCombinedDivisors(3, CombinedDivisorPattern.ThreeAOneB, out CombinedDivisorsEnding3ThreeAOneBSquares);
-		CombinedDivisorsEnding7ThreeAOneB = BuildCombinedDivisors(7, CombinedDivisorPattern.ThreeAOneB);
+		CombinedDivisorsEnding7ThreeAOneB = BuildCombinedDivisors(7, CombinedDivisorPattern.ThreeAOneB, out CombinedDivisorsEnding7ThreeAOneBSquares);
-		CombinedDivisorsEnding9ThreeAOneB = BuildCombinedDivisors(9, CombinedDivisorPattern.ThreeAOneB);
+		CombinedDivisorsEnding9ThreeAOneB = BuildCombinedDivisors(9, CombinedDivisorPattern.ThreeAOneB, out CombinedDivisorsEnding9ThreeAOneBSquares);
 
-		CombinedDivisorsEnding1ThreeATwoB = BuildCombinedDivisors(1, CombinedDivisorPattern.ThreeATwoB);
+		CombinedDivisorsEnding1ThreeATwoB = BuildCombinedDivisors(1, CombinedDivisorPattern.ThreeATwoB, out CombinedDivisorsEnding1ThreeATwoBSquares);
-		CombinedDivisorsEnding3ThreeATwoB = BuildCombinedDivisors(3, CombinedDivisorPattern.ThreeATwoB);
+		CombinedDivisorsEnding3ThreeATwoB = BuildCombinedDivisors(3, CombinedDivisorPattern.ThreeATwoB, out CombinedDivisorsEnding3ThreeATwoBSquares);
-		CombinedDivisorsEnding7ThreeATwoB = BuildCombinedDivisors(7, CombinedDivisorPattern.ThreeATwoB);
+		CombinedDivisorsEnding7ThreeATwoB = BuildCombinedDivisors(7, CombinedDivisorPattern.ThreeATwoB, out CombinedDivisorsEnding7ThreeATwoBSquares);
-		CombinedDivisorsEnding9ThreeATwoB = BuildCombinedDivisors(9, CombinedDivisorPattern.ThreeATwoB);
+		CombinedDivisorsEnding9ThreeATwoB = BuildCombinedDivisors(9, CombinedDivisorPattern.ThreeATwoB, out CombinedDivisorsEnding9ThreeATwoBSquares);
 		CombinedDivisorsEnding1 = CombinedDivisorsEnding1OneAOneB;
 		CombinedDivisorsEnding3 = CombinedDivisorsEnding3OneAOneB;
 		CombinedDivisorsEnding7 = CombinedDivisorsEnding7OneAOneB;
 		CombinedDivisorsEnding9 = CombinedDivisorsEnding9OneAOneB;
+		CombinedDivisorsEnding1Squares = CombinedDivisorsEnding1OneAOneBSquares;
+		CombinedDivisorsEnding3Squares = CombinedDivisorsEnding3OneAOneBSquares;
+		CombinedDivisorsEnding7Squares = CombinedDivisorsEnding7OneAOneBSquares;
+		CombinedDivisorsEnding9Squares = CombinedDivisorsEnding9OneAOneBSquares;
 	}
 
 	private const ulong Wheel210 = 210UL;
@@ -71,6 +75,31 @@ public sealed class HeuristicCombinedPrimeTester
 	private static readonly uint[] CombinedDivisorsEnding7;
 	private static readonly uint[] CombinedDivisorsEnding9;
 
+	private static readonly ulong[] CombinedDivisorsEnding1TwoAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding3TwoAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding7TwoAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding9TwoAOneBSquares;
+
+	private static readonly ulong[] CombinedDivisorsEnding1OneAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding3OneAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding7OneAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding9OneAOneBSquares;
+
+	private static readonly ulong[] CombinedDivisorsEnding1ThreeAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding3ThreeAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding7ThreeAOneBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding9ThreeAOneBSquares;
+
+	private static readonly ulong[] CombinedDivisorsEnding1ThreeATwoBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding3ThreeATwoBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding7ThreeATwoBSquares;
+	private static readonly ulong[] CombinedDivisorsEnding9ThreeATwoBSquares;
+
+	private static readonly ulong[] CombinedDivisorsEnding1Squares;
+	private static readonly ulong[] CombinedDivisorsEnding3Squares;
+	private static readonly ulong[] CombinedDivisorsEnding7Squares;
+	private static readonly ulong[] CombinedDivisorsEnding9Squares;
+
 	private static readonly ulong[] HeuristicSmallCycleSnapshot = MersenneDivisorCycles.Shared.ExportSmallCyclesSnapshot();
 
 	internal enum CombinedDivisorPattern : byte
@@ -155,21 +184,28 @@ public sealed class HeuristicCombinedPrimeTester
 	private static bool HeuristicTrialDivisionCpu(ulong n, ulong sqrtLimit, byte nMod10)
 	{
 		ReadOnlySpan<uint> combinedDivisors = GetCombinedDivisors(nMod10);
-		if (combinedDivisors.IsEmpty)
+		ReadOnlySpan<ulong> combinedDivisorSquares = GetCombinedDivisorSquares(nMod10);
+		if (combinedDivisors.IsEmpty || combinedDivisorSquares.IsEmpty)
 		{
 			return EvaluateWithOpenNumericFallback(n);
 		}
 
 		int length = combinedDivisors.Length;
+		int squaresLength = combinedDivisorSquares.Length;
+		if (squaresLength < length)
+		{
+			length = squaresLength;
+		}
+
 		for (int i = 0; i < length; i++)
 		{
-			uint entry = combinedDivisors[i];
+			ulong divisorSquare = combinedDivisorSquares[i];
-			if (entry > sqrtLimit)
+			if (divisorSquare > n)
 			{
 				break;
 			}
 
-			if (n % entry == 0UL)
+			if (n % combinedDivisors[i] == 0UL)
 			{
 				return false;
 			}
@@ -182,64 +218,57 @@ public sealed class HeuristicCombinedPrimeTester
 	{
 	        int batchCapacity = HeuristicGpuDivisorBatchSize;
 	        var divisorPool = ThreadStaticPools.UlongPool;
-		var hitPool = ThreadStaticPools.BytePool;
 
 	        ulong[] divisorArray = divisorPool.Rent(batchCapacity);
-		byte[] hitFlags = hitPool.Rent(batchCapacity);
 
 	        var limiter = GpuPrimeWorkLimiter.Acquire();
-		var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent();
+	        var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity, 1);
 	        var accelerator = gpu.Accelerator;
-		var state = PrimeTester.GpuKernelState.GetOrCreate(accelerator);
+	        var state = gpu.State;
-		var scratch = state.RentScratch(batchCapacity, accelerator);
+	        var scratch = gpu.Scratch;
 
 	        bool compositeDetected = false;
 	        int count = 0;
 
-		var input = scratch.Input;
+	        var inputView = scratch.Input.View;
-		var output = scratch.Output;
+	        var outputView = scratch.Output.View;
 
 	        bool ProcessBatch(int length)
 	        {
-			input.View.CopyFromCPU(ref divisorArray[0], length);
+	                inputView.CopyFromCPU(ref divisorArray[0], length);
-			state.HeuristicTrialDivisionKernel(length, input.View, n, output.View);
+	                byte compositeFlag = 0;
+	                outputView.CopyFromCPU(ref compositeFlag, 1);
+	                state.HeuristicTrialDivisionKernel(length, inputView, n, outputView);
 	                accelerator.Synchronize();
-			output.View.CopyToCPU(ref hitFlags[0], length);
+	                outputView.CopyToCPU(ref compositeFlag, 1);
 
-			for (int i = 0; i < length; i++)
+	                return compositeFlag != 0;
-			{
-				if (hitFlags[i] == 0)
-				{
-					continue;
-				}
-
-				return true;
-				// ulong divisor = divisorArray[i];
-				// if (divisor > 1UL && n % divisor == 0UL)
-				// {
-				// 	return true;
-				// }
-			}
-
-			return false;
 	        }
 
 	        ReadOnlySpan<uint> combinedDivisors = GetCombinedDivisors(nMod10);
-		if (combinedDivisors.IsEmpty)
+                ReadOnlySpan<ulong> combinedDivisorSquares = GetCombinedDivisorSquares(nMod10);
+
+                if (combinedDivisors.IsEmpty || combinedDivisorSquares.IsEmpty)
                 {
                         goto Cleanup;
                 }
 
                 int length = combinedDivisors.Length;
+                int squaresLength = combinedDivisorSquares.Length;
+                if (squaresLength < length)
+                {
+                        length = squaresLength;
+                }
+
                 for (int i = 0; i < length; i++)
                 {
-			uint entry = combinedDivisors[i];
+                        ulong divisorSquare = combinedDivisorSquares[i];
-			if (entry > sqrtLimit)
+                        if (divisorSquare > n)
                         {
                                 break;
                         }
 
-			divisorArray[count] = entry;
+                        divisorArray[count] = combinedDivisors[i];
                         count++;
 
                         if (count == batchCapacity)
@@ -254,24 +283,21 @@ public sealed class HeuristicCombinedPrimeTester
                         }
                 }
 
-		if (!compositeDetected && count > 0)
+                if (count > 0 && ProcessBatch(count))
-		{
-			if (ProcessBatch(count))
                 {
                         compositeDetected = true;
                         goto Cleanup;
                 }
-		}
+
 
 Cleanup:
-		state.ReturnScratch(scratch);
 	        divisorPool.Return(divisorArray);
-		hitPool.Return(hitFlags);
 	        gpu.Dispose();
 	        limiter.Dispose();
 	        return compositeDetected;
 	}
 
+
 	private static bool EvaluateWithOpenNumericFallback(ulong n)
 	{
 		return Prime.Numbers.IsPrime(n);
@@ -312,6 +338,18 @@ public sealed class HeuristicCombinedPrimeTester
 		};
 	}
 
+	private static ReadOnlySpan<ulong> GetCombinedDivisorSquares(byte nMod10)
+	{
+		return nMod10 switch
+		{
+			1 => CombinedDivisorsEnding1Squares,
+			3 => CombinedDivisorsEnding3Squares,
+			7 => CombinedDivisorsEnding7Squares,
+			9 => CombinedDivisorsEnding9Squares,
+			_ => ReadOnlySpan<ulong>.Empty,
+		};
+	}
+
 	internal static ReadOnlySpan<uint> GetCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern)
 	{
 		return pattern switch
@@ -352,7 +390,47 @@ public sealed class HeuristicCombinedPrimeTester
 		};
 	}
 
-	private static uint[] BuildCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern)
+	internal static ReadOnlySpan<ulong> GetCombinedDivisorSquares(byte nMod10, CombinedDivisorPattern pattern)
+	{
+		return pattern switch
+		{
+			CombinedDivisorPattern.TwoAOneB => nMod10 switch
+			{
+				1 => CombinedDivisorsEnding1TwoAOneBSquares,
+				3 => CombinedDivisorsEnding3TwoAOneBSquares,
+				7 => CombinedDivisorsEnding7TwoAOneBSquares,
+				9 => CombinedDivisorsEnding9TwoAOneBSquares,
+				_ => ReadOnlySpan<ulong>.Empty,
+			},
+			CombinedDivisorPattern.OneAOneB => nMod10 switch
+			{
+				1 => CombinedDivisorsEnding1OneAOneBSquares,
+				3 => CombinedDivisorsEnding3OneAOneBSquares,
+				7 => CombinedDivisorsEnding7OneAOneBSquares,
+				9 => CombinedDivisorsEnding9OneAOneBSquares,
+				_ => ReadOnlySpan<ulong>.Empty,
+			},
+			CombinedDivisorPattern.ThreeAOneB => nMod10 switch
+			{
+				1 => CombinedDivisorsEnding1ThreeAOneBSquares,
+				3 => CombinedDivisorsEnding3ThreeAOneBSquares,
+				7 => CombinedDivisorsEnding7ThreeAOneBSquares,
+				9 => CombinedDivisorsEnding9ThreeAOneBSquares,
+				_ => ReadOnlySpan<ulong>.Empty,
+			},
+			CombinedDivisorPattern.ThreeATwoB => nMod10 switch
+			{
+				1 => CombinedDivisorsEnding1ThreeATwoBSquares,
+				3 => CombinedDivisorsEnding3ThreeATwoBSquares,
+				7 => CombinedDivisorsEnding7ThreeATwoBSquares,
+				9 => CombinedDivisorsEnding9ThreeATwoBSquares,
+				_ => ReadOnlySpan<ulong>.Empty,
+			},
+			_ => ReadOnlySpan<ulong>.Empty,
+		};
+	}
+
+	private static uint[] BuildCombinedDivisors(byte nMod10, CombinedDivisorPattern pattern, out ulong[] squares)
 	{
 		ReadOnlySpan<int> groupA = HeuristicPrimeSieves.GroupADivisors;
 		ReadOnlySpan<uint> groupB = GetGroupBDivisors(nMod10);
@@ -412,7 +490,16 @@ public sealed class HeuristicCombinedPrimeTester
 			}
 		}
 
-		return combined.ToArray();
+		uint[] result = combined.ToArray();
+		ulong[] resultSquares = new ulong[result.Length];
+		for (int i = 0; i < result.Length; i++)
+		{
+			ulong divisor = result[i];
+			resultSquares[i] = divisor * divisor;
+		}
+
+		squares = resultSquares;
+		return result;
 	}
 
 

PerfectNumbers.Core/HeuristicPrimeTester.cs

@@ -263,52 +263,31 @@ public sealed class HeuristicPrimeTester
         int batchCapacity = HeuristicGpuDivisorBatchSize;
         int interleaveBatchSize = HeuristicDivisorInterleaveBatchSize;
         var divisorPool = ThreadStaticPools.UlongPool;
-        var hitPool = ThreadStaticPools.BytePool;
 
         ulong[] divisorArray = divisorPool.Rent(batchCapacity);
-        byte[] hitFlags = hitPool.Rent(batchCapacity);
 
         var limiter = GpuPrimeWorkLimiter.Acquire();
-        var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent();
+        var gpu = PrimeTester.PrimeTesterGpuContextPool.Rent(batchCapacity, 1);
         var accelerator = gpu.Accelerator;
-        var state = PrimeTester.GpuKernelState.GetOrCreate(accelerator);
+        var state = gpu.State;
-        var scratch = state.RentScratch(batchCapacity, accelerator);
+        var scratch = gpu.Scratch;
 
         bool compositeDetected = false;
         int count = 0;
 
-        // lock (gpu.ExecutionLock)
+        var inputView = scratch.Input.View;
-        {
+        var outputView = scratch.Output.View;
-            var input = scratch.Input;
-            var output = scratch.Output;
 
         bool ProcessBatch(int length)
         {
-                if (length <= 0)
+            inputView.CopyFromCPU(ref divisorArray[0], length);
-                {
+            byte compositeFlag = 0;
-                    return false;
+            outputView.CopyFromCPU(ref compositeFlag, 1);
-                }
+            state.HeuristicTrialDivisionKernel(length, inputView, n, outputView);
-
-                input.View.CopyFromCPU(ref divisorArray[0], length);
-                state.HeuristicTrialDivisionKernel(length, input.View, n, output.View);
             accelerator.Synchronize();
-                output.View.CopyToCPU(ref hitFlags[0], length);
+            outputView.CopyToCPU(ref compositeFlag, 1);
 
-                for (int i = 0; i < length; i++)
+            return compositeFlag != 0;
-                {
-                    if (hitFlags[i] == 0)
-                    {
-                        continue;
-                    }
-
-                    ulong divisor = divisorArray[i];
-                    if (divisor > 1UL && n % divisor == 0UL)
-                    {
-                        return true;
-                    }
-                }
-
-                return false;
         }
 
         ReadOnlySpan<int> groupADivisors = HeuristicPrimeSieves.GroupADivisors;
@@ -403,27 +382,20 @@ public sealed class HeuristicPrimeTester
             }
         }
 
-            if (!compositeDetected && count > 0)
+        if (count > 0 && ProcessBatch(count))
-            {
-                if (ProcessBatch(count))
         {
             compositeDetected = true;
             goto Cleanup;
         }
 
-                count = 0;
-            }
-        }
-
 Cleanup:
-        state.ReturnScratch(scratch);
         divisorPool.Return(divisorArray);
-        hitPool.Return(hitFlags);
         gpu.Dispose();
         limiter.Dispose();
         return compositeDetected;
     }
 
+
     private static bool EvaluateWithOpenNumericFallback(ulong n)
     {
         return Prime.Numbers.IsPrime(n);

PerfectNumbers.Core/PrimeTester.cs

@@ -128,13 +128,13 @@ public sealed class PrimeTester
 	public static void IsPrimeBatchGpu(ReadOnlySpan<ulong> values, Span<byte> results)
 	{
 		var limiter = GpuPrimeWorkLimiter.Acquire();
-		var gpu = PrimeTesterGpuContextPool.Rent();
+		var gpu = PrimeTesterGpuContextPool.Rent(GpuBatchSize);
+		var scratch = gpu.Scratch;
 		var state = gpu.State;
 		var accelerator = gpu.Accelerator;
 		int totalLength = values.Length;
-		int batchSize = Math.Max(1, GpuBatchSize);
+		int batchSize = GpuBatchSize;
 
-		var scratch = state.RentScratch(batchSize, accelerator);
 		var input = scratch.Input;
 		var output = scratch.Output;
 		ArrayPool<ulong> pool = ThreadStaticPools.UlongPool;
@@ -170,7 +170,6 @@ public sealed class PrimeTester
 		}
 
 		pool.Return(temp, clearArray: false);
-		state.ReturnScratch(scratch);
 
 		gpu.Dispose();
 		limiter.Dispose();
@@ -185,6 +184,7 @@ public sealed class PrimeTester
 			public Accelerator Accelerator { get; }
 
 			private KernelState? _kernelState;
+			private KernelState.ScratchBuffers? _scratch;
 
 			public PooledContext()
 			{
@@ -208,8 +208,51 @@ public sealed class PrimeTester
 				}
 			}
 
+			public KernelState.ScratchBuffers GetScratch(int minCapacity, int outputClearLength)
+			{
+				if (minCapacity <= 0)
+				{
+					minCapacity = 1;
+				}
+
+				var scratch = _scratch;
+				var state = KernelState;
+				if (scratch is null)
+				{
+					scratch = state.RentScratch(minCapacity, Accelerator);
+					_scratch = scratch;
+				}
+				else if (scratch.Capacity < minCapacity)
+				{
+					scratch.Dispose();
+					scratch = state.RentScratch(minCapacity, Accelerator);
+					_scratch = scratch;
+				}
+
+				int clearLength = outputClearLength <= 0 ? 0 : outputClearLength;
+				if (clearLength > 0)
+				{
+					int outputLength = (int)scratch.Output.Length;
+					if (clearLength > outputLength)
+					{
+						clearLength = outputLength;
+					}
+
+					scratch.Output.View.SubView(0, clearLength).MemSetToZero();
+					Accelerator.Synchronize();
+				}
+
+				return scratch;
+			}
+
 			public void Dispose()
 			{
+				if (_scratch is not null)
+				{
+					_scratch.Dispose();
+					_scratch = null;
+				}
+
 				ClearGpuCaches(Accelerator);
 				_kernelState = null;
 				Accelerator.Dispose();
@@ -219,14 +262,14 @@ public sealed class PrimeTester
 
 		private static readonly ConcurrentQueue<PooledContext> Pool = new();
 
-		internal static PrimeTesterGpuContextLease Rent()
+		internal static PrimeTesterGpuContextLease Rent(int minScratchCapacity = 0, int outputClearLength = 0)
 		{
 			if (Pool.TryDequeue(out var ctx))
 			{
-				return new PrimeTesterGpuContextLease(ctx);
+				return new PrimeTesterGpuContextLease(ctx, minScratchCapacity, outputClearLength);
 			}
 
-			return new PrimeTesterGpuContextLease(new PooledContext());
+			return new PrimeTesterGpuContextLease(new PooledContext(), minScratchCapacity, outputClearLength);
 		}
 
 		internal static void DisposeAll()
@@ -243,18 +286,30 @@ public sealed class PrimeTester
 			Pool.Enqueue(ctx);
 		}
 
-		internal readonly struct PrimeTesterGpuContextLease(PrimeTesterGpuContextPool.PooledContext ctx)
+		internal readonly struct PrimeTesterGpuContextLease
 		{
-			private readonly PooledContext? _ctx = ctx;
+			private readonly PooledContext? _ctx;
+			private readonly KernelState? _state;
+			private readonly KernelState.ScratchBuffers? _scratch;
 
-			public readonly Accelerator Accelerator => _ctx!.Accelerator;
+			internal PrimeTesterGpuContextLease(PooledContext ctx, int minScratchCapacity, int outputClearLength)
-
-
-			public readonly KernelState State => _ctx!.KernelState;
-
-			public readonly void Dispose()
 			{
-				Return(_ctx ?? throw new InvalidOperationException("GPU context lease is not initialized."));
+				_ctx = ctx;
+				_state = ctx.KernelState;
+				_scratch = minScratchCapacity > 0 ? ctx.GetScratch(minScratchCapacity, outputClearLength) : null;
+			}
+
+			public Accelerator Accelerator => (_ctx ?? throw new InvalidOperationException("GPU context lease is not initialized.")).Accelerator;
+
+			public KernelState State => _state ?? throw new InvalidOperationException("GPU kernel state is not initialized.");
+
+			public KernelState.ScratchBuffers Scratch => _scratch ?? throw new InvalidOperationException("Scratch buffers were not rented for this lease.");
+
+			public void Dispose()
+			{
+				var ctx = _ctx ?? throw new InvalidOperationException("GPU context lease is not initialized.");
+
+				Return(ctx);
 			}
 		}
 	}