Reuse GPU context for nested partial factoring

This commit is contained in:
Marek Lemańczyk committed 2025-10-26 20:58:21 +01:00
1 parent 6458f6740d
commit ba439f0f0f
4 files changed
+270 -146

No files matched your search

+14 -2
View File
@@ -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!);
}
}
@@ -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)
+113 -21
View File
@@ -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,
+27 -41
View File
@@ -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
{