Reuse GPU context for nested partial factoring

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -46,6 +46,7 @@ public sealed class KernelContainer
     public MemoryBuffer1D<int, Stride1D.Dense>? SmallPrimeFactorExponentSlots;
     public MemoryBuffer1D<int, Stride1D.Dense>? SmallPrimeFactorCountSlot;
     public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimeFactorRemainingSlot;
+    public MemoryBuffer1D<byte, Stride1D.Dense>? SmallPrimeFactorFullyFactoredSlot;
     public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastOne;
     public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimesPow2LastOne;
     public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastSeven;
@@ -107,17 +108,20 @@ public readonly struct SmallPrimeFactorScratch
     public readonly MemoryBuffer1D<int, Stride1D.Dense> ExponentSlots;
     public readonly MemoryBuffer1D<int, Stride1D.Dense> CountSlot;
     public readonly MemoryBuffer1D<ulong, Stride1D.Dense> RemainingSlot;
+    public readonly MemoryBuffer1D<byte, Stride1D.Dense> FullyFactoredSlot;
 
     public SmallPrimeFactorScratch(
         MemoryBuffer1D<ulong, Stride1D.Dense> primeSlots,
         MemoryBuffer1D<int, Stride1D.Dense> exponentSlots,
         MemoryBuffer1D<int, Stride1D.Dense> countSlot,
-        MemoryBuffer1D<ulong, Stride1D.Dense> remainingSlot)
+        MemoryBuffer1D<ulong, Stride1D.Dense> remainingSlot,
+        MemoryBuffer1D<byte, Stride1D.Dense> fullyFactoredSlot)
     {
         PrimeSlots = primeSlots;
         ExponentSlots = exponentSlots;
         CountSlot = countSlot;
         RemainingSlot = remainingSlot;
+        FullyFactoredSlot = fullyFactoredSlot;
     }
 
     public void Clear()
@@ -126,6 +130,7 @@ public readonly struct SmallPrimeFactorScratch
         ExponentSlots.MemSetToZero();
         CountSlot.MemSetToZero();
         RemainingSlot.MemSetToZero();
+        FullyFactoredSlot.MemSetToZero();
     }
 }
 
@@ -273,7 +278,14 @@ public class GpuKernelPool
                 kernels.SmallPrimeFactorRemainingSlot = remainingSlot;
             }
 
-            return new SmallPrimeFactorScratch(primeSlots!, exponentSlots!, countSlot!, remainingSlot!);
+            MemoryBuffer1D<byte, Stride1D.Dense>? fullyFactoredSlot = kernels.SmallPrimeFactorFullyFactoredSlot;
+            if (fullyFactoredSlot is null)
+            {
+                fullyFactoredSlot = accelerator.Allocate1D<byte>(1);
+                kernels.SmallPrimeFactorFullyFactoredSlot = fullyFactoredSlot;
+            }
+
+            return new SmallPrimeFactorScratch(primeSlots!, exponentSlots!, countSlot!, remainingSlot!, fullyFactoredSlot!);
         }
     }
 

PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs

@@ -151,34 +151,52 @@ internal static partial class PrimeOrderGpuHeuristics
         primeTargets.Clear();
         exponentTargets.Clear();
 
-        if (!TryLaunchPartialFactorKernel(
-                value,
-                limit,
-                primeTargets,
-                exponentTargets,
-                out int extracted,
-                out ulong leftover,
-                out bool kernelFullyFactored))
+        try
         {
-            return false;
+            var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+            var execution = lease.EnterExecutionScope();
+            try
+            {
+                Accelerator accelerator = lease.Accelerator;
+                AcceleratorStream stream = lease.Stream;
+                SmallPrimeFactorScratch scratch = GpuKernelPool.EnsureSmallPrimeFactorScratch(accelerator, primeTargets.Length);
+                return TryPartialFactor(
+                    accelerator,
+                    stream,
+                    scratch,
+                    value,
+                    limit,
+                    primeTargets,
+                    exponentTargets,
+                    out factorCount,
+                    out remaining,
+                    out fullyFactored);
+            }
+            finally
+            {
+                execution.Dispose();
+                lease.Dispose();
+            }
+        }
+        catch (CLException ex)
+        {
+            Console.WriteLine($"GPU ERROR ({ex.Error}): {ex.Message}");
+        }
+        catch (Exception ex) when (ex is AcceleratorException or InternalCompilerException or NotSupportedException or InvalidOperationException or AggregateException)
+        {
+            Console.WriteLine($"GPU ERROR: {ex.Message}");
         }
 
-        int capacity = Math.Min(primeTargets.Length, exponentTargets.Length);
-        if (extracted > capacity)
-        {
-            factorCount = 0;
-            remaining = value;
-            fullyFactored = false;
-            return false;
-        }
-
-        factorCount = extracted;
-        remaining = leftover;
-        fullyFactored = kernelFullyFactored && leftover == 1UL;
-        return true;
+        factorCount = 0;
+        remaining = value;
+        fullyFactored = false;
+        return false;
     }
 
-    private static bool TryLaunchPartialFactorKernel(
+    public static bool TryPartialFactor(
+        Accelerator accelerator,
+        AcceleratorStream stream,
+        SmallPrimeFactorScratch scratch,
         ulong value,
         uint limit,
         Span<ulong> primeTargets,
@@ -191,83 +209,99 @@ internal static partial class PrimeOrderGpuHeuristics
         remaining = value;
         fullyFactored = false;
 
+        primeTargets.Clear();
+        exponentTargets.Clear();
+
         try
         {
-            var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
-            var execution = lease.EnterExecutionScope();
-            Accelerator accelerator = lease.Accelerator;
-            AcceleratorStream stream = lease.Stream;
-
-            var kernel = GetPartialFactorKernel(accelerator);
-			SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator);
-
-			// TODO: We should create / reallocate these buffer only once or if the new required length is bigger than capacity. 
-            var factorBuffer = accelerator.Allocate1D<ulong>(primeTargets.Length);
-            var exponentBuffer = accelerator.Allocate1D<int>(exponentTargets.Length);
-            var countBuffer = accelerator.Allocate1D<int>(1);
-            var remainingBuffer = accelerator.Allocate1D<ulong>(1);
-            var fullyFactoredBuffer = accelerator.Allocate1D<byte>(1);
-
-			// TODO: There is no need to clear these buffers because the kernel will always assign values within the required bounds.
-			// factorBuffer.MemSetToZero();
-            // exponentBuffer.MemSetToZero();
-            // countBuffer.MemSetToZero();
-            // remainingBuffer.MemSetToZero();
-            // fullyFactoredBuffer.MemSetToZero();
-
-            kernel(
+            return TryPartialFactorCore(
+                accelerator,
                 stream,
-                1,
-                cache.Primes!.View,
-                cache.Squares!.View,
-                cache.Count,
-                primeTargets.Length,
+                scratch,
                 value,
                 limit,
-                factorBuffer.View,
-                exponentBuffer.View,
-                countBuffer.View,
-                remainingBuffer.View,
-                fullyFactoredBuffer.View);
-
-            stream.Synchronize();
-
-            countBuffer.View.CopyToCPU(ref factorCount, 1);
-            factorCount = Math.Min(factorCount, primeTargets.Length);
-            factorBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(primeTargets), primeTargets.Length);
-            exponentBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(exponentTargets), exponentTargets.Length);
-            remainingBuffer.View.CopyToCPU(ref remaining, 1);
-
-            byte fullyFactoredFlag = 0;
-            fullyFactoredBuffer.View.CopyToCPU(ref fullyFactoredFlag, 1);
-            fullyFactored = fullyFactoredFlag != 0;
-
-			// TODO: These buffers shouldn't be disposed but rather reused accross call with the assigned accelerator.
-            factorBuffer.Dispose();
-            exponentBuffer.Dispose();
-            countBuffer.Dispose();
-            remainingBuffer.Dispose();
-            fullyFactoredBuffer.Dispose();
-            execution.Dispose();
-            lease.Dispose();
-            return true;
+                primeTargets,
+                exponentTargets,
+                out factorCount,
+                out remaining,
+                out fullyFactored);
         }
         catch (CLException ex)
         {
             Console.WriteLine($"GPU ERROR ({ex.Error}): {ex.Message}");
-            factorCount = 0;
-            remaining = value;
-            fullyFactored = false;
-            return false;
         }
         catch (Exception ex) when (ex is AcceleratorException or InternalCompilerException or NotSupportedException or InvalidOperationException or AggregateException)
         {
             Console.WriteLine($"GPU ERROR: {ex.Message}");
+        }
+
+        factorCount = 0;
+        remaining = value;
+        fullyFactored = false;
+        return false;
+    }
+
+    private static bool TryPartialFactorCore(
+        Accelerator accelerator,
+        AcceleratorStream stream,
+        SmallPrimeFactorScratch scratch,
+        ulong value,
+        uint limit,
+        Span<ulong> primeTargets,
+        Span<int> exponentTargets,
+        out int factorCount,
+        out ulong remaining,
+        out bool fullyFactored)
+    {
+        factorCount = 0;
+        remaining = value;
+        fullyFactored = false;
+
+        scratch.Clear();
+
+        var kernel = GetPartialFactorKernel(accelerator);
+        SmallPrimeDeviceCache cache = GetSmallPrimeDeviceCache(accelerator);
+
+        kernel(
+            stream,
+            1,
+            cache.Primes!.View,
+            cache.Squares!.View,
+            cache.Count,
+            primeTargets.Length,
+            value,
+            limit,
+            scratch.PrimeSlots.View,
+            scratch.ExponentSlots.View,
+            scratch.CountSlot.View,
+            scratch.RemainingSlot.View,
+            scratch.FullyFactoredSlot.View);
+
+        stream.Synchronize();
+
+        scratch.CountSlot.View.CopyToCPU(ref factorCount, 1);
+        int capacity = Math.Min(primeTargets.Length, exponentTargets.Length);
+        if (factorCount > capacity)
+        {
             factorCount = 0;
             remaining = value;
             fullyFactored = false;
             return false;
         }
+
+        if (factorCount > 0)
+        {
+            scratch.PrimeSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(primeTargets), factorCount);
+            scratch.ExponentSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(exponentTargets), factorCount);
+        }
+
+        scratch.RemainingSlot.View.CopyToCPU(ref remaining, 1);
+
+        byte fullyFactoredFlag = 0;
+        scratch.FullyFactoredSlot.View.CopyToCPU(ref fullyFactoredFlag, 1);
+        bool kernelFullyFactored = fullyFactoredFlag != 0;
+        fullyFactored = kernelFullyFactored && remaining == 1UL;
+        return true;
     }
 
     public static GpuPow2ModStatus TryPow2Mod(ulong exponent, ulong prime, out ulong remainder, in MontgomeryDivisorData divisorData)