From ba439f0f0fba496a02e0133645710d9acc76eacb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Marek=20Lema=C5=84czyk?= Date: Sun, 26 Oct 2025 20:58:21 +0100 Subject: [PATCH] Reuse GPU context for nested partial factoring --- PerfectNumbers.Core/Gpu/GpuKernelPool.cs | 16 +- .../Gpu/PrimeOrderGpuHeuristics.cs | 198 ++++++++++-------- .../PrimeOrderCalculator.Cpu.cs | 134 ++++++++++-- .../PrimeOrderCalculator.Gpu.cs | 68 +++--- 4 files changed, 270 insertions(+), 146 deletions(-) diff --git a/PerfectNumbers.Core/Gpu/GpuKernelPool.cs b/PerfectNumbers.Core/Gpu/GpuKernelPool.cs index c16d9ca..2060b4f 100644 --- a/PerfectNumbers.Core/Gpu/GpuKernelPool.cs +++ b/PerfectNumbers.Core/Gpu/GpuKernelPool.cs @@ -46,6 +46,7 @@ public sealed class KernelContainer public MemoryBuffer1D? SmallPrimeFactorExponentSlots; public MemoryBuffer1D? SmallPrimeFactorCountSlot; public MemoryBuffer1D? SmallPrimeFactorRemainingSlot; + public MemoryBuffer1D? SmallPrimeFactorFullyFactoredSlot; public MemoryBuffer1D? SmallPrimesLastOne; public MemoryBuffer1D? SmallPrimesPow2LastOne; public MemoryBuffer1D? SmallPrimesLastSeven; @@ -107,17 +108,20 @@ public readonly struct SmallPrimeFactorScratch public readonly MemoryBuffer1D ExponentSlots; public readonly MemoryBuffer1D CountSlot; public readonly MemoryBuffer1D RemainingSlot; + public readonly MemoryBuffer1D FullyFactoredSlot; public SmallPrimeFactorScratch( MemoryBuffer1D primeSlots, MemoryBuffer1D exponentSlots, MemoryBuffer1D countSlot, - MemoryBuffer1D remainingSlot) + MemoryBuffer1D remainingSlot, + MemoryBuffer1D 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? fullyFactoredSlot = kernels.SmallPrimeFactorFullyFactoredSlot; + if (fullyFactoredSlot is null) + { + fullyFactoredSlot = accelerator.Allocate1D(1); + kernels.SmallPrimeFactorFullyFactoredSlot = fullyFactoredSlot; + } + + return new SmallPrimeFactorScratch(primeSlots!, exponentSlots!, countSlot!, remainingSlot!, fullyFactoredSlot!); } } diff --git a/PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs b/PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs index 7b17b48..04ed0ce 100644 --- a/PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs +++ b/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 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(primeTargets.Length); - var exponentBuffer = accelerator.Allocate1D(exponentTargets.Length); - var countBuffer = accelerator.Allocate1D(1); - var remainingBuffer = accelerator.Allocate1D(1); - var fullyFactoredBuffer = accelerator.Allocate1D(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 primeTargets, + Span 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) diff --git a/PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs b/PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs index 9b83a9c..70e8b6c 100644 --- a/PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs +++ b/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,19 +96,38 @@ internal static partial class PrimeOrderCalculator return gpuOrder; } - PartialFactorResult phiFactors = PartialFactor(phi, config); + PartialFactorGpuContext? gpuContext = null; + bool disposeContext = false; - ulong result; - if (phiFactors.Factors is null) + if (IsGpuHeuristicDevice) { - result = CalculateByFactorizationCpu(prime, divisorData); + gpuContext = PartialFactorGpuContext.Rent(); + disposeContext = true; + } + + try + { + PartialFactorResult phiFactors = PartialFactor(phi, config, gpuContext); + + ulong result; + if (phiFactors.Factors is null) + { + result = CalculateByFactorizationCpu(prime, divisorData); + phiFactors.Dispose(); + return result; + } + + result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors, gpuContext); phiFactors.Dispose(); return result; } - - result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors); - phiFactors.Dispose(); - return result; + finally + { + if (disposeContext && gpuContext is not null) + { + gpuContext.Dispose(); + } + } } private static ulong RunHeuristicPipelineCpu( @@ -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, diff --git a/PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs b/PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs index 4cb801c..8e6caf8 100644 --- a/PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs +++ b/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,52 +31,39 @@ 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; - kernel( - stream, - 1, - value, - limit, - tables.PrimesView, - tables.SquaresView, - tables.Count, - scratch.PrimeSlots.View, - scratch.ExponentSlots.View, - scratch.CountSlot.View, - scratch.RemainingSlot.View); + bool success = PrimeOrderGpuHeuristics.TryPartialFactor( + accelerator, + stream, + scratch, + value, + limit, + primeBuffer, + exponentBuffer, + out factorCount, + out remaining, + out _); - stream.Synchronize(); + if (!success) + { + return false; + } - factorCount = 0; - scratch.CountSlot.View.CopyToCPU(ref factorCount, 1); - factorCount = Math.Min(factorCount, GpuSmallPrimeFactorSlots); + for (int i = 0; i < factorCount; i++) + { + ulong primeValue = primeBuffer[i]; + int exponent = exponentBuffer[i]; + // This will never happen in production code + // if (primeValue == 0UL || exponent == 0) + // { + // continue; + // } - if (factorCount > 0) - { - scratch.PrimeSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(primeBuffer), factorCount); - scratch.ExponentSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(exponentBuffer), factorCount); - } + counts.Add(primeValue, exponent); + } - scratch.RemainingSlot.View.CopyToCPU(ref remaining, 1); - - for (int i = 0; i < factorCount; i++) - { - ulong primeValue = primeBuffer[i]; - int exponent = exponentBuffer[i]; - // This will never happen in production code - // if (primeValue == 0UL || exponent == 0) - // { - // continue; - // } - - counts.Add(primeValue, exponent); - } - - return true; + return true; } finally {