Merge pull request #86 from mlemanczyk/codex/implement-todo-for-cpu-prime-residue-calculations

Move cycle filtering and residue prep to CPU
This commit is contained in:
Marek Lemańczyk authored and GitHub committed 2025-09-25 06:10:03 +02:00
commit 8c3e9cb76f
3 files changed
+10 -78

No files matched your search

+3 -26
View File
@@ -882,10 +882,6 @@ internal static class Program
int hitIndex = 0; int hitIndex = 0;
int index; int index;
bool useDivisorCycles = _byDivisorTester!.UseDivisorCycles; bool useDivisorCycles = _byDivisorTester!.UseDivisorCycles;
DivisorCycleCache.Lease cycleLease = default;
ulong cycleLeaseStart = 0UL;
ulong cycleLeaseEnd = 0UL;
ulong[]? cycleValues = null;
ulong divisorCycle = 0UL; ulong divisorCycle = 0UL;
try try
@@ -903,27 +899,9 @@ internal static class Program
divisor = AcquireNextDivisor(ref nextDivisor, divisorLimit, ref divisorsExhaustedFlag, ref finalDivisorBits, out exhausted, ref localDivisorCursor, ref localDivisorsRemaining); divisor = AcquireNextDivisor(ref nextDivisor, divisorLimit, ref divisorsExhaustedFlag, ref finalDivisorBits, out exhausted, ref localDivisorCursor, ref localDivisorsRemaining);
if (divisor != 0UL) if (divisor != 0UL)
{ {
if (useDivisorCycles) divisorCycle = useDivisorCycles
{ ? DivisorCycleCache.Shared.Acquire(divisor).GetCycle(divisor)
if (!cycleLease.IsValid || divisor < cycleLeaseStart || divisor > cycleLeaseEnd) : 0UL;
{
cycleLease.Dispose();
cycleLease = DivisorCycleCache.Shared.Acquire(divisor);
cycleLeaseStart = cycleLease.Start;
cycleLeaseEnd = cycleLease.End;
cycleValues = cycleLease.Values ?? throw new InvalidOperationException("Divisor cycle cache returned no values for acquired lease.");
}
else if (cycleValues is null)
{
cycleValues = cycleLease.Values ?? throw new InvalidOperationException("Divisor cycle cache returned no values for acquired lease.");
}
divisorCycle = cycleValues[(int)(divisor - cycleLeaseStart)];
}
else
{
divisorCycle = 0UL;
}
activeCount = BuildPrimeBuffer(divisor, primeValues, allowedMaxValues, stateFlags, primeBuffer, indexBuffer, completionsBuffer, ref completionsCount, ref remainingStates, activeStateMask, ref activeStartIndex); activeCount = BuildPrimeBuffer(divisor, primeValues, allowedMaxValues, stateFlags, primeBuffer, indexBuffer, completionsBuffer, ref completionsCount, ref remainingStates, activeStateMask, ref activeStartIndex);
@@ -1018,7 +996,6 @@ internal static class Program
} }
finally finally
{ {
cycleLease.Dispose();
ArrayPool<PendingResult>.Shared.Return(compositesBuffer, clearArray: true); ArrayPool<PendingResult>.Shared.Return(compositesBuffer, clearArray: true);
ArrayPool<PendingResult>.Shared.Return(completionsBuffer, clearArray: true); ArrayPool<PendingResult>.Shared.Return(completionsBuffer, clearArray: true);
ArrayPool<int>.Shared.Return(indexBuffer, clearArray: true); ArrayPool<int>.Shared.Return(indexBuffer, clearArray: true);
@@ -171,8 +171,8 @@ public sealed class MersenneNumberDivisorByDivisorCpuTester : IMersenneNumberDiv
byte hit; byte hit;
if (useCycles) if (useCycles)
{ {
using DivisorCycleCache.Lease cycleLease = DivisorCycleCache.Shared.Acquire(divisor); DivisorCycleCache.CycleBlock cycleBlock = DivisorCycleCache.Shared.Acquire(divisor);
ulong divisorCycle = cycleLease.GetCycle(divisor); ulong divisorCycle = cycleBlock.GetCycle(divisor);
hit = CheckDivisor(prime, divisorCycle != 0UL, divisorCycle, divisorData); hit = CheckDivisor(prime, divisorCycle != 0UL, divisorCycle, divisorData);
} }
else else
@@ -26,15 +26,12 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ulong, byte, ulong, ArrayView<MontgomeryDivisorData>, ArrayView<byte>>> _kernelCache = new(); private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ulong, byte, ulong, ArrayView<MontgomeryDivisorData>, ArrayView<byte>>> _kernelCache = new();
private readonly ConcurrentDictionary<Accelerator, Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>>> _kernelExponentCache = new(); private readonly ConcurrentDictionary<Accelerator, Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>>> _kernelExponentCache = new();
private readonly ConcurrentDictionary<Accelerator, ConcurrentBag<BatchResources>> _resourcePools = new(AcceleratorReferenceComparer.Instance); private readonly ConcurrentDictionary<Accelerator, ConcurrentBag<BatchResources>> _resourcePools = new(AcceleratorReferenceComparer.Instance);
private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>>> _cycleKernelCache = new();
private readonly ConcurrentBag<DivisorScanSession> _sessionPool = new(); private readonly ConcurrentBag<DivisorScanSession> _sessionPool = new();
private Action<Index1D, ulong, byte, ulong, ArrayView<MontgomeryDivisorData>, ArrayView<byte>> GetKernel(Accelerator accelerator) => private Action<Index1D, ulong, byte, ulong, ArrayView<MontgomeryDivisorData>, ArrayView<byte>> GetKernel(Accelerator accelerator) =>
_kernelCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, ulong, byte, ulong, ArrayView<MontgomeryDivisorData>, ArrayView<byte>>(CheckKernel)); _kernelCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, ulong, byte, ulong, ArrayView<MontgomeryDivisorData>, ArrayView<byte>>(CheckKernel));
private Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> GetKernelByPrimeExponent(Accelerator accelerator) => private Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> GetKernelByPrimeExponent(Accelerator accelerator) =>
_kernelExponentCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>>(ComputePrimeExponentKernel)); _kernelExponentCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>>(ComputePrimeExponentKernel));
private Action<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>> GetCycleKernel(Accelerator accelerator) =>
_cycleKernelCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>>(CheckPrimeCycleKernel));
public int GpuBatchSize public int GpuBatchSize
{ {
@@ -334,14 +331,10 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
private readonly GpuContextPool.GpuContextLease _lease; private readonly GpuContextPool.GpuContextLease _lease;
private readonly Accelerator _accelerator; private readonly Accelerator _accelerator;
private readonly Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> _kernel; private readonly Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> _kernel;
private readonly Action<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>> _cycleKernel;
private MemoryBuffer1D<ulong, Stride1D.Dense> _exponentsBuffer; private MemoryBuffer1D<ulong, Stride1D.Dense> _exponentsBuffer;
private MemoryBuffer1D<ulong, Stride1D.Dense> _resultsBuffer; private MemoryBuffer1D<ulong, Stride1D.Dense> _resultsBuffer;
private MemoryBuffer1D<byte, Stride1D.Dense> _cycleMaskBuffer;
private ulong[] _hostBuffer; private ulong[] _hostBuffer;
private ulong[] _primeBuffer;
private int[] _positionBuffer; private int[] _positionBuffer;
private byte[] _cycleMaskHost;
private int _capacity; private int _capacity;
private bool _disposed; private bool _disposed;
@@ -351,15 +344,11 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
_lease = GpuContextPool.RentPreferred(preferCpu: false); _lease = GpuContextPool.RentPreferred(preferCpu: false);
_accelerator = _lease.Accelerator; _accelerator = _lease.Accelerator;
_kernel = owner.GetKernelByPrimeExponent(_accelerator); _kernel = owner.GetKernelByPrimeExponent(_accelerator);
_cycleKernel = owner.GetCycleKernel(_accelerator);
_capacity = Math.Max(1, owner._gpuBatchSize); _capacity = Math.Max(1, owner._gpuBatchSize);
_exponentsBuffer = _accelerator.Allocate1D<ulong>(_capacity); _exponentsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
_resultsBuffer = _accelerator.Allocate1D<ulong>(_capacity); _resultsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
_cycleMaskBuffer = _accelerator.Allocate1D<byte>(_capacity);
_hostBuffer = ArrayPool<ulong>.Shared.Rent(_capacity); _hostBuffer = ArrayPool<ulong>.Shared.Rent(_capacity);
_primeBuffer = ArrayPool<ulong>.Shared.Rent(_capacity);
_positionBuffer = ArrayPool<int>.Shared.Rent(_capacity); _positionBuffer = ArrayPool<int>.Shared.Rent(_capacity);
_cycleMaskHost = ArrayPool<byte>.Shared.Rent(_capacity);
} }
internal void Reset() internal void Reset()
@@ -376,20 +365,14 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
_exponentsBuffer.Dispose(); _exponentsBuffer.Dispose();
_resultsBuffer.Dispose(); _resultsBuffer.Dispose();
_cycleMaskBuffer.Dispose();
ArrayPool<ulong>.Shared.Return(_hostBuffer, clearArray: false); ArrayPool<ulong>.Shared.Return(_hostBuffer, clearArray: false);
ArrayPool<ulong>.Shared.Return(_primeBuffer, clearArray: false);
ArrayPool<int>.Shared.Return(_positionBuffer, clearArray: false); ArrayPool<int>.Shared.Return(_positionBuffer, clearArray: false);
ArrayPool<byte>.Shared.Return(_cycleMaskHost, clearArray: false);
_capacity = requiredCapacity; _capacity = requiredCapacity;
_exponentsBuffer = _accelerator.Allocate1D<ulong>(_capacity); _exponentsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
_resultsBuffer = _accelerator.Allocate1D<ulong>(_capacity); _resultsBuffer = _accelerator.Allocate1D<ulong>(_capacity);
_cycleMaskBuffer = _accelerator.Allocate1D<byte>(_capacity);
_hostBuffer = ArrayPool<ulong>.Shared.Rent(_capacity); _hostBuffer = ArrayPool<ulong>.Shared.Rent(_capacity);
_primeBuffer = ArrayPool<ulong>.Shared.Rent(_capacity);
_positionBuffer = ArrayPool<int>.Shared.Rent(_capacity); _positionBuffer = ArrayPool<int>.Shared.Rent(_capacity);
_cycleMaskHost = ArrayPool<byte>.Shared.Rent(_capacity);
} }
public void CheckDivisor(ulong divisor, ulong divisorCycle, ReadOnlySpan<ulong> primes, Span<byte> hits) public void CheckDivisor(ulong divisor, ulong divisorCycle, ReadOnlySpan<ulong> primes, Span<byte> hits)
@@ -415,13 +398,10 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
Accelerator accelerator = _accelerator; Accelerator accelerator = _accelerator;
Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> kernel = _kernel; Action<Index1D, MontgomeryDivisorData, ArrayView<ulong>, ArrayView<ulong>> kernel = _kernel;
Action<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>> cycleKernel = _cycleKernel;
ArrayView1D<ulong, Stride1D.Dense> exponentsView = _exponentsBuffer.View; ArrayView1D<ulong, Stride1D.Dense> exponentsView = _exponentsBuffer.View;
ArrayView1D<ulong, Stride1D.Dense> resultsView = _resultsBuffer.View; ArrayView1D<ulong, Stride1D.Dense> resultsView = _resultsBuffer.View;
ArrayView1D<byte, Stride1D.Dense> cycleMaskViewFull = _cycleMaskBuffer.View;
Span<ulong> hostSpan = _hostBuffer.AsSpan(); Span<ulong> hostSpan = _hostBuffer.AsSpan();
Span<ulong> primeSpan = _primeBuffer.AsSpan();
Span<int> positionSpan = _positionBuffer.AsSpan(); Span<int> positionSpan = _positionBuffer.AsSpan();
int offset = 0; int offset = 0;
@@ -434,26 +414,18 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
int computeCount = 0; int computeCount = 0;
if (cycleEnabled) if (cycleEnabled && divisorCycle != 0UL)
{ {
Span<byte> cycleMaskSpan = _cycleMaskHost.AsSpan(0, batchSize);
ArrayView1D<ulong, Stride1D.Dense> primeInputView = _resultsBuffer.View.SubView(0, batchSize);
primeInputView.CopyFromCPU(ref MemoryMarshal.GetReference(primesSlice), batchSize);
ArrayView1D<byte, Stride1D.Dense> cycleMaskView = cycleMaskViewFull.SubView(0, batchSize);
// TODO: Move this back to CPU side and make sure that we only pass primes here which passed the cycle test, when cycles are in use.
cycleKernel(batchSize, divisorCycle, primeInputView, cycleMaskView);
accelerator.Synchronize();
cycleMaskView.CopyToCPU(ref MemoryMarshal.GetReference(cycleMaskSpan), batchSize);
for (int i = 0; i < batchSize; i++) for (int i = 0; i < batchSize; i++)
{ {
if (cycleMaskSpan[i] == 0) ulong residue = primesSlice[i] % divisorCycle;
if (residue != 0UL)
{ {
hitsSlice[i] = 0; hitsSlice[i] = 0;
continue; continue;
} }
primeSpan[computeCount] = primesSlice[i]; hostSpan[computeCount] = residue;
positionSpan[computeCount] = i; positionSpan[computeCount] = i;
computeCount++; computeCount++;
} }
@@ -462,7 +434,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
{ {
for (int i = 0; i < batchSize; i++) for (int i = 0; i < batchSize; i++)
{ {
primeSpan[computeCount] = primesSlice[i]; hostSpan[computeCount] = primesSlice[i];
positionSpan[computeCount] = i; positionSpan[computeCount] = i;
computeCount++; computeCount++;
} }
@@ -471,12 +443,6 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
if (computeCount > 0) if (computeCount > 0)
{ {
Span<ulong> exponentSlice = hostSpan[..computeCount]; Span<ulong> exponentSlice = hostSpan[..computeCount];
for (int i = 0; i < computeCount; i++)
{
ulong prime = primeSpan[i];
// TODO: Let's calculate prime % divisorCycle on CPU side, too. When we check the cycle, we calculate the residue. Let's put them into an array and re-use here.
exponentSlice[i] = cycleEnabled ? prime % divisorCycle : prime;
}
ArrayView1D<ulong, Stride1D.Dense> exponentView = exponentsView.SubView(0, computeCount); ArrayView1D<ulong, Stride1D.Dense> exponentView = exponentsView.SubView(0, computeCount);
exponentView.CopyFromCPU(ref MemoryMarshal.GetReference(exponentSlice), computeCount); exponentView.CopyFromCPU(ref MemoryMarshal.GetReference(exponentSlice), computeCount);
@@ -516,17 +482,6 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
_sessionPool.Add(session); _sessionPool.Add(session);
} }
private static void CheckPrimeCycleKernel(Index1D index, ulong cycle, ArrayView<ulong> primes, ArrayView<byte> mask)
{
if (cycle == 0UL)
{
mask[index] = 1;
return;
}
mask[index] = primes[index] % cycle == 0UL ? (byte)1 : (byte)0;
}
private static void ComputePrimeExponentKernel(Index1D index, MontgomeryDivisorData divisor, ArrayView<ulong> exponents, ArrayView<ulong> results) private static void ComputePrimeExponentKernel(Index1D index, MontgomeryDivisorData divisor, ArrayView<ulong> exponents, ArrayView<ulong> results)
{ {
ulong modulus = divisor.Modulus; ulong modulus = divisor.Modulus;