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(
+        try
+        {
+            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 int extracted,
+                    out factorCount,
-                out ulong leftover,
+                    out remaining,
-                out bool kernelFullyFactored))
+                    out fullyFactored);
+            }
+            finally
             {
-            return false;
+                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;
+    public static bool TryPartialFactor(
-        remaining = leftover;
+        Accelerator accelerator,
-        fullyFactored = kernelFullyFactored && leftover == 1UL;
+        AcceleratorStream stream,
-        return true;
+        SmallPrimeFactorScratch scratch,
-    }
-
-    private static bool TryLaunchPartialFactorKernel(
         ulong value,
         uint limit,
         Span<ulong> primeTargets,
@@ -191,30 +209,59 @@ internal static partial class PrimeOrderGpuHeuristics
         remaining = value;
         fullyFactored = false;
 
+        primeTargets.Clear();
+        exponentTargets.Clear();
+
         try
         {
-            var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+            return TryPartialFactorCore(
-            var execution = lease.EnterExecutionScope();
+                accelerator,
-            Accelerator accelerator = lease.Accelerator;
+                stream,
-            AcceleratorStream stream = lease.Stream;
+                scratch,
+                value,
+                limit,
+                primeTargets,
+                exponentTargets,
+                out factorCount,
+                out remaining,
+                out fullyFactored);
+        }
+        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}");
+        }
+
+        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);
 
-			// 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(
             stream,
             1,
@@ -224,51 +271,38 @@ internal static partial class PrimeOrderGpuHeuristics
             primeTargets.Length,
             value,
             limit,
-                factorBuffer.View,
+            scratch.PrimeSlots.View,
-                exponentBuffer.View,
+            scratch.ExponentSlots.View,
-                countBuffer.View,
+            scratch.CountSlot.View,
-                remainingBuffer.View,
+            scratch.RemainingSlot.View,
-                fullyFactoredBuffer.View);
+            scratch.FullyFactoredSlot.View);
 
         stream.Synchronize();
 
-            countBuffer.View.CopyToCPU(ref factorCount, 1);
+        scratch.CountSlot.View.CopyToCPU(ref factorCount, 1);
-            factorCount = Math.Min(factorCount, primeTargets.Length);
+        int capacity = Math.Min(primeTargets.Length, exponentTargets.Length);
-            factorBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(primeTargets), primeTargets.Length);
+        if (factorCount > capacity)
-            exponentBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(exponentTargets), exponentTargets.Length);
+        {
-            remainingBuffer.View.CopyToCPU(ref remaining, 1);
+            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;
-            fullyFactoredBuffer.View.CopyToCPU(ref fullyFactoredFlag, 1);
+        scratch.FullyFactoredSlot.View.CopyToCPU(ref fullyFactoredFlag, 1);
-            fullyFactored = fullyFactoredFlag != 0;
+        bool kernelFullyFactored = fullyFactoredFlag != 0;
-
+        fullyFactored = kernelFullyFactored && remaining == 1UL;
-			// 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;
     }
-        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;
-        }
-    }
 
     public static GpuPow2ModStatus TryPow2Mod(ulong exponent, ulong prime, out ulong remainder, in MontgomeryDivisorData divisorData)
     {

PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs

@@ -3,9 +3,10 @@ using System.Buffers;
 using System.Diagnostics;
 using System.Numerics;
 using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 using ILGPU;
+using ILGPU.Runtime;
 using PerfectNumbers.Core.Gpu;
+using System.Runtime.InteropServices;
 
 namespace PerfectNumbers.Core;
 
@@ -43,6 +44,43 @@ internal static partial class PrimeOrderCalculator
 		}
 	}
 
+        private sealed class PartialFactorGpuContext : IDisposable
+        {
+                private readonly GpuKernelLease _lease;
+                private readonly GpuKernelLease.ExecutionScope _executionScope;
+
+                private PartialFactorGpuContext(GpuKernelLease lease, GpuKernelLease.ExecutionScope executionScope, Accelerator accelerator, AcceleratorStream stream, SmallPrimeFactorScratch scratch)
+                {
+                        _lease = lease;
+                        _executionScope = executionScope;
+                        Accelerator = accelerator;
+                        Stream = stream;
+                        Scratch = scratch;
+                }
+
+                public Accelerator Accelerator { get; }
+
+                public AcceleratorStream Stream { get; }
+
+                public SmallPrimeFactorScratch Scratch { get; }
+
+                public static PartialFactorGpuContext Rent()
+                {
+                        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+                        var execution = lease.EnterExecutionScope();
+                        Accelerator accelerator = lease.Accelerator;
+                        AcceleratorStream stream = lease.Stream;
+                        SmallPrimeFactorScratch scratch = GpuKernelPool.EnsureSmallPrimeFactorScratch(accelerator, GpuSmallPrimeFactorSlots);
+                        return new PartialFactorGpuContext(lease, execution, accelerator, stream, scratch);
+                }
+
+                public void Dispose()
+                {
+                        _executionScope.Dispose();
+                        _lease.Dispose();
+                }
+        }
+
 	private static ulong CalculateInternal(ulong prime, ulong? previousOrder, in MontgomeryDivisorData divisorData, in PrimeOrderSearchConfig config)
 	{
 		// TODO: Is this condition ever met on EvenPerfectBitScanner's execution path? If not, we can add a clarification comment and comment out the entire block. We want to support p candidates at least greater or equal to 31.
@@ -58,7 +96,18 @@ internal static partial class PrimeOrderCalculator
 			return gpuOrder;
 		}
 
-		PartialFactorResult phiFactors = PartialFactor(phi, config);
+		PartialFactorGpuContext? gpuContext = null;
+		bool disposeContext = false;
+
+		if (IsGpuHeuristicDevice)
+		{
+			gpuContext = PartialFactorGpuContext.Rent();
+			disposeContext = true;
+		}
+
+		try
+		{
+			PartialFactorResult phiFactors = PartialFactor(phi, config, gpuContext);
 
 			ulong result;
 			if (phiFactors.Factors is null)
@@ -68,10 +117,18 @@ internal static partial class PrimeOrderCalculator
 				return result;
 			}
 
-		result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors);
+			result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors, gpuContext);
 			phiFactors.Dispose();
 			return result;
 		}
+		finally
+		{
+			if (disposeContext && gpuContext is not null)
+			{
+				gpuContext.Dispose();
+			}
+		}
+	}
 
 	private static ulong RunHeuristicPipelineCpu(
 		ulong prime,
@@ -79,7 +136,8 @@ internal static partial class PrimeOrderCalculator
 		in PrimeOrderSearchConfig config,
 		in MontgomeryDivisorData divisorData,
 		ulong phi,
-		PartialFactorResult phiFactors)
+		PartialFactorResult phiFactors,
+		PartialFactorGpuContext? gpuContext)
 	{
 		if (phiFactors.FullyFactored && TrySpecialMaxCpu(phi, prime, phiFactors, divisorData))
 		{
@@ -89,7 +147,7 @@ internal static partial class PrimeOrderCalculator
 		ulong candidateOrder = InitializeStartingOrderCpu(prime, phi, divisorData);
 		candidateOrder = ExponentLoweringCpu(candidateOrder, prime, phiFactors, divisorData);
 
-		if (TryConfirmOrderCpu(prime, candidateOrder, divisorData, config, out PartialFactorResult? orderFactors))
+		if (TryConfirmOrderCpu(prime, candidateOrder, divisorData, config, gpuContext, out PartialFactorResult? orderFactors))
 		{
 			return candidateOrder;
 		}
@@ -100,7 +158,7 @@ internal static partial class PrimeOrderCalculator
 			return CalculateByFactorizationCpu(prime, divisorData);
 		}
 
-		if (TryHeuristicFinishCpu(prime, candidateOrder, previousOrder, divisorData, config, phiFactors, orderFactors, out ulong order))
+		if (TryHeuristicFinishCpu(prime, candidateOrder, previousOrder, divisorData, config, phiFactors, orderFactors, gpuContext, out ulong order))
 		{
 			return order;
 		}
@@ -330,6 +388,7 @@ internal static partial class PrimeOrderCalculator
 		ulong order,
 		in MontgomeryDivisorData divisorData,
 		in PrimeOrderSearchConfig config,
+		PartialFactorGpuContext? gpuContext,
 		out PartialFactorResult? reusableFactorization)
 	{
 		reusableFactorization = null;
@@ -344,7 +403,7 @@ internal static partial class PrimeOrderCalculator
 			return false;
 		}
 
-		PartialFactorResult factorization = PartialFactor(order, config);
+		PartialFactorResult factorization = PartialFactor(order, config, gpuContext);
 		if (factorization.Factors is null)
 		{
 			factorization.Dispose();
@@ -492,6 +551,7 @@ internal static partial class PrimeOrderCalculator
 		in PrimeOrderSearchConfig config,
 		PartialFactorResult phiFactors,
 		PartialFactorResult? cachedOrderFactors,
+		PartialFactorGpuContext? gpuContext,
 		out ulong result)
 	{
 		result = 0UL;
@@ -502,7 +562,7 @@ internal static partial class PrimeOrderCalculator
 		}
 
 		// Reuse the partial factorization from TryConfirmOrderCpu when available.
-		PartialFactorResult orderFactors = cachedOrderFactors ?? PartialFactor(order, config);
+		PartialFactorResult orderFactors = cachedOrderFactors ?? PartialFactor(order, config, gpuContext);
 		try
 		{
 			if (orderFactors.Factors is null)
@@ -618,7 +678,7 @@ internal static partial class PrimeOrderCalculator
 
 					bool equalsOne = powStepper.InitializeCpuIsUnity(candidate);
 					powStepperInitialized = true;
-					if (equalsOne && TryConfirmCandidateCpu(prime, candidate, divisorData, config, ref powUsed, powBudget))
+					if (equalsOne && TryConfirmCandidateCpu(prime, candidate, divisorData, config, gpuContext, ref powUsed, powBudget))
 					{
 						if (gpuPool is not null)
 						{
@@ -655,7 +715,7 @@ internal static partial class PrimeOrderCalculator
 						continue;
 					}
 
-					if (!TryConfirmCandidateCpu(prime, candidate, divisorData, config, ref powUsed, powBudget))
+					if (!TryConfirmCandidateCpu(prime, candidate, divisorData, config, gpuContext, ref powUsed, powBudget))
 					{
 						continue;
 					}
@@ -858,9 +918,9 @@ internal static partial class PrimeOrderCalculator
 		}
 	}
 
-	private static bool TryConfirmCandidateCpu(ulong prime, ulong candidate, in MontgomeryDivisorData divisorData, in PrimeOrderSearchConfig config, ref int powUsed, int powBudget)
+	private static bool TryConfirmCandidateCpu(ulong prime, ulong candidate, in MontgomeryDivisorData divisorData, in PrimeOrderSearchConfig config, PartialFactorGpuContext? gpuContext, ref int powUsed, int powBudget)
 	{
-		PartialFactorResult factorization = PartialFactor(candidate, config);
+		PartialFactorResult factorization = PartialFactor(candidate, config, gpuContext);
 		try
 		{
 			if (factorization.Factors is null)
@@ -1024,7 +1084,7 @@ internal static partial class PrimeOrderCalculator
 	// private static ulong _partialFactorPendingHits;
 	// private static ulong _partialFactorCofactorHits;
 
-	private static PartialFactorResult PartialFactor(ulong value, in PrimeOrderSearchConfig config)
+	private static PartialFactorResult PartialFactor(ulong value, in PrimeOrderSearchConfig config, PartialFactorGpuContext? gpuContext = null)
 	{
 		if (value <= 1UL)
 		{
@@ -1033,6 +1093,17 @@ internal static partial class PrimeOrderCalculator
 
 		const int FactorSlotCount = GpuSmallPrimeFactorSlots;
 
+                PartialFactorGpuContext? localContext = gpuContext;
+                bool disposeContext = false;
+
+                if (IsGpuHeuristicDevice && localContext is null)
+                {
+                        localContext = PartialFactorGpuContext.Rent();
+                        disposeContext = true;
+                }
+
+                try
+                {
 		// stackalloc is faster than pooling
 		// ulong[] primeSlotsArray = ThreadStaticPools.UlongPool.Rent(FactorSlotCount);
 		// int[] exponentSlotsArray = ThreadStaticPools.IntPool.Rent(FactorSlotCount);
@@ -1052,7 +1123,21 @@ internal static partial class PrimeOrderCalculator
 		ulong remaining = value;
 
 		bool gpuFactored = false;
-		if (IsGpuHeuristicDevice)
+		if (IsGpuHeuristicDevice && localContext is not null)
+		{
+			gpuFactored = PrimeOrderGpuHeuristics.TryPartialFactor(
+				localContext.Accelerator,
+				localContext.Stream,
+				localContext.Scratch,
+				value,
+				limit,
+				primeSlots,
+				exponentSlots,
+				out factorCount,
+				out remaining,
+				out _);
+		}
+		else if (IsGpuHeuristicDevice)
 		{
 			gpuFactored = PrimeOrderGpuHeuristics.TryPartialFactor(value, limit, primeSlots, exponentSlots, out factorCount, out remaining, out _);
 		}
@@ -1329,7 +1414,14 @@ internal static partial class PrimeOrderCalculator
 		ThreadStaticPools.ReturnPrimeOrderPendingEntryList(pending);
 		return result;
 	}
-
+	finally
+	{
+		if (disposeContext && localContext is not null)
+		{
+			localContext.Dispose();
+		}
+	}
+	}
 
 	private static bool TryPopulateSmallPrimeFactorsCpu(
 			ulong value,

PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs

@@ -1,5 +1,4 @@
 using System;
-using System.Runtime.InteropServices;
 using ILGPU;
 using ILGPU.Runtime;
 using ILGPU.Runtime.OpenCL;
@@ -32,38 +31,25 @@ internal static partial class PrimeOrderCalculator
 				Accelerator accelerator = lease.Accelerator;
 				AcceleratorStream stream = lease.Stream;
 
-				SmallPrimeFactorTables tables = GpuKernelPool.EnsureSmallPrimeFactorTables(accelerator);
 				SmallPrimeFactorScratch scratch = GpuKernelPool.EnsureSmallPrimeFactorScratch(accelerator, GpuSmallPrimeFactorSlots);
-				// scratch.Clear();
 
-				var kernel = lease.SmallPrimeFactorKernel;
+                                bool success = PrimeOrderGpuHeuristics.TryPartialFactor(
-				kernel(
+                                        accelerator,
                                         stream,
-					1,
+                                        scratch,
                                         value,
                                         limit,
-					tables.PrimesView,
+                                        primeBuffer,
-					tables.SquaresView,
+                                        exponentBuffer,
-					tables.Count,
+                                        out factorCount,
-					scratch.PrimeSlots.View,
+                                        out remaining,
-					scratch.ExponentSlots.View,
+                                        out _);
-					scratch.CountSlot.View,
-					scratch.RemainingSlot.View);
 
-				stream.Synchronize();
+                                if (!success)
-
-				factorCount = 0;
-				scratch.CountSlot.View.CopyToCPU(ref factorCount, 1);
-				factorCount = Math.Min(factorCount, GpuSmallPrimeFactorSlots);
-
-				if (factorCount > 0)
                                 {
-					scratch.PrimeSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(primeBuffer), factorCount);
+                                        return false;
-					scratch.ExponentSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(exponentBuffer), factorCount);
                                 }
 
-				scratch.RemainingSlot.View.CopyToCPU(ref remaining, 1);
-
                                 for (int i = 0; i < factorCount; i++)
                                 {
                                         ulong primeValue = primeBuffer[i];