Merge pull request #57 from mlemanczyk/codex/optimize-memory-usage-for-mersenne-method

Stream the by-divisor GPU tester without retaining divisors

PerfectNumbers.Core.Tests/MersenneNumberDivisorGpuTesterTests.cs

@@ -102,5 +102,6 @@ public class MersenneNumberDivisorGpuTesterTests
         tester.IsPrime(11UL, out bool divisorsExhausted).Should().BeTrue();
         divisorsExhausted.Should().BeTrue();
     }
+
 }
 

PerfectNumbers.Core/Gpu/MersenneNumberDivisorByDivisorGpuTester.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Buffers;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
 using System.Runtime.InteropServices;
 using ILGPU;
 using ILGPU.Algorithms;
@@ -12,23 +11,15 @@ namespace PerfectNumbers.Core.Gpu;
 public sealed class MersenneNumberDivisorByDivisorGpuTester
 {
         private int _gpuBatchSize = GpuConstants.ScanBatchSize;
-
-        private readonly List<ulong> _divisors = new();
-        private readonly List<ulong> _remainders = new();
-        private readonly List<ulong> _lastPrimes = new();
         private readonly object _sync = new();
         private ulong _divisorLimit;
         private bool _isConfigured;
         private ulong _lastStatusDivisor;
 
-        private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ulong, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>> _updateKernelCache = new();
+        private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>>> _kernelCache = new();
-        private readonly ConcurrentDictionary<Accelerator, Action<Index1D, ulong, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>> _initialKernelCache = new();
 
-        private Action<Index1D, ulong, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> GetUpdateKernel(Accelerator accelerator) =>
+        private Action<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>> GetKernel(Accelerator accelerator) =>
-                        _updateKernelCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<ulong>, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(UpdateKernel));
+                        _kernelCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<ulong>, ArrayView<byte>>(CheckKernel));
-
-        private Action<Index1D, ulong, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>> GetInitialKernel(Accelerator accelerator) =>
-                        _initialKernelCache.GetOrAdd(accelerator, acc => acc.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView<ulong>, ArrayView<ulong>, ArrayView<byte>>(InitialKernel));
 
         public int GpuBatchSize
         {
@@ -40,13 +31,6 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester
         {
                 lock (_sync)
                 {
-                        if (_divisors.Count > 0)
-                        {
-                                _divisors.Clear();
-                                _remainders.Clear();
-                                _lastPrimes.Clear();
-                        }
-
                         _divisorLimit = ComputeDivisorLimitFromMaxPrime(maxPrime);
                         _lastStatusDivisor = 0UL;
                         _isConfigured = true;
@@ -69,192 +53,103 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester
                                 return true;
                         }
 
-                        bool composite = UpdateExistingDivisors(prime, allowedMax);
+                        bool composite = CheckDivisors(prime, allowedMax, out ulong lastProcessed, out bool coveredRange);
-                        if (!composite)
+
+                        if (!composite && lastProcessed != 0UL)
                         {
-                                composite = ExtendDivisors(prime, allowedMax);
+                                ReportStatus(lastProcessed);
                         }
 
-                        divisorsExhausted = true;
+                        if (composite)
-                        return !composite;
+                        {
+                                divisorsExhausted = true;
+                                return false;
+                        }
+
+                        divisorsExhausted = coveredRange;
+                        return true;
                 }
         }
 
-        private bool UpdateExistingDivisors(ulong prime, ulong allowedMax)
+        private bool CheckDivisors(ulong prime, ulong allowedMax, out ulong lastProcessed, out bool coveredRange)
         {
-                int count = _divisors.Count;
-                if (count == 0)
-                {
-                        return false;
-                }
-
-                var divisorsSpan = CollectionsMarshal.AsSpan(_divisors);
-                var remaindersSpan = CollectionsMarshal.AsSpan(_remainders);
-                var lastPrimesSpan = CollectionsMarshal.AsSpan(_lastPrimes);
-
-                var gpuLease = GpuContextPool.RentPreferred(preferCpu: false);
-                var accelerator = gpuLease.Accelerator;
-                var kernel = GetUpdateKernel(accelerator);
-
-                int batchCapacity = Math.Min(_gpuBatchSize, count);
-                var divisorsBuffer = accelerator.Allocate1D<ulong>(batchCapacity);
-                var remaindersBuffer = accelerator.Allocate1D<ulong>(batchCapacity);
-                var lastPrimesBuffer = accelerator.Allocate1D<ulong>(batchCapacity);
-                var hitsBuffer = accelerator.Allocate1D<byte>(batchCapacity);
-
-                byte[] hits = ArrayPool<byte>.Shared.Rent(batchCapacity);
-
-                bool composite = false;
-                ulong lastProcessed = 0UL;
-
-                try
-                {
-                        int processed = 0;
-                        while (processed < count)
-                        {
-                                int batchSize = Math.Min(batchCapacity, count - processed);
-                                var divisorView = divisorsBuffer.View.SubView(0, batchSize);
-                                var remainderView = remaindersBuffer.View.SubView(0, batchSize);
-                                var lastPrimeView = lastPrimesBuffer.View.SubView(0, batchSize);
-                                var hitsView = hitsBuffer.View.SubView(0, batchSize);
-
-                                Span<ulong> divisorBatch = divisorsSpan.Slice(processed, batchSize);
-                                Span<ulong> remainderBatch = remaindersSpan.Slice(processed, batchSize);
-                                Span<ulong> lastPrimeBatch = lastPrimesSpan.Slice(processed, batchSize);
-
-                                divisorView.CopyFromCPU(ref MemoryMarshal.GetReference(divisorBatch), batchSize);
-                                remainderView.CopyFromCPU(ref MemoryMarshal.GetReference(remainderBatch), batchSize);
-                                lastPrimeView.CopyFromCPU(ref MemoryMarshal.GetReference(lastPrimeBatch), batchSize);
-                                hitsView.MemSetToZero();
-
-                                kernel(batchSize, prime, divisorView, remainderView, lastPrimeView, hitsView);
-                                accelerator.Synchronize();
-
-                                remainderView.CopyToCPU(ref MemoryMarshal.GetReference(remainderBatch), batchSize);
-                                lastPrimeView.CopyToCPU(ref MemoryMarshal.GetReference(lastPrimeBatch), batchSize);
-
-                                Span<byte> hitsSpan = hits.AsSpan(0, batchSize);
-                                hitsView.CopyToCPU(ref MemoryMarshal.GetReference(hitsSpan), batchSize);
-
-                                for (int i = 0; i < batchSize; i++)
-                                {
-                                        ulong divisor = divisorBatch[i];
-                                        if (divisor > allowedMax)
-                                        {
-                                                processed = count;
-                                                break;
-                                        }
-
-                                        lastProcessed = divisor;
-                                        if (hitsSpan[i] != 0)
-                                        {
-                                                composite = true;
-                                                processed = count;
-                                                break;
-                                        }
-                                }
-
-                                if (composite)
-                                {
-                                        break;
-                                }
-
-                                if (processed < count)
-                                {
-                                        processed += batchSize;
-                                }
-                        }
-                }
-                finally
-                {
-                        ArrayPool<byte>.Shared.Return(hits, clearArray: true);
-                        hitsBuffer.Dispose();
-                        lastPrimesBuffer.Dispose();
-                        remaindersBuffer.Dispose();
-                        divisorsBuffer.Dispose();
-                        gpuLease.Dispose();
-                }
-
-                if (lastProcessed != 0UL)
-                {
-                        ReportStatus(lastProcessed);
-                }
-
-                return composite;
-        }
-
-        private bool ExtendDivisors(ulong prime, ulong allowedMax)
-        {
-                int existingCount = _divisors.Count;
                 allowedMax = Math.Min(allowedMax, _divisorLimit);
+                lastProcessed = 0UL;
 
-                ulong currentMax = existingCount == 0 ? 1UL : _divisors[^1];
+                if (allowedMax < 3UL)
-                ulong start = currentMax < 3UL ? 3UL : currentMax + 2UL;
-                if (start > allowedMax)
                 {
+                        coveredRange = true;
                         return false;
                 }
 
                 var gpuLease = GpuContextPool.RentPreferred(preferCpu: false);
                 var accelerator = gpuLease.Accelerator;
-                var kernel = GetInitialKernel(accelerator);
+                var kernel = GetKernel(accelerator);
 
                 int batchCapacity = _gpuBatchSize;
                 var divisorsBuffer = accelerator.Allocate1D<ulong>(batchCapacity);
-                var remaindersBuffer = accelerator.Allocate1D<ulong>(batchCapacity);
                 var hitsBuffer = accelerator.Allocate1D<byte>(batchCapacity);
 
-                ulong[] remainders = ArrayPool<ulong>.Shared.Rent(batchCapacity);
-                byte[] hits = ArrayPool<byte>.Shared.Rent(batchCapacity);
                 ulong[] divisors = ArrayPool<ulong>.Shared.Rent(batchCapacity);
+                byte[] hits = ArrayPool<byte>.Shared.Rent(batchCapacity);
 
                 bool composite = false;
-                ulong lastProcessed = 0UL;
+                bool processedAll = false;
-                ulong divisor = start;
+                ulong divisor = 3UL;
 
                 try
                 {
                         while (divisor <= allowedMax)
                         {
-                                int batchSize = batchCapacity;
+                                int batchSize = 0;
-                                int index = 0;
+                                bool reachedEndInBatch = false;
-                                while (index < batchSize && divisor <= allowedMax)
+
+                                while (batchSize < batchCapacity && divisor <= allowedMax)
                                 {
-                                        divisors[index++] = divisor;
+                                        ulong current = divisor;
-                                        divisor += 2UL;
+                                        divisors[batchSize++] = current;
+
+                                        ulong next = current + 2UL;
+                                        if (next <= current)
+                                        {
+                                                reachedEndInBatch = true;
+                                                break;
+                                        }
+
+                                        if (next > allowedMax)
+                                        {
+                                                divisor = next;
+                                                reachedEndInBatch = true;
+                                                break;
+                                        }
+
+                                        divisor = next;
                                 }
 
-                                if (index == 0)
+                                if (batchSize == 0)
                                 {
                                         break;
                                 }
 
-                                var divisorView = divisorsBuffer.View.SubView(0, index);
+                                var divisorView = divisorsBuffer.View.SubView(0, batchSize);
-                                var remainderView = remaindersBuffer.View.SubView(0, index);
+                                var hitsView = hitsBuffer.View.SubView(0, batchSize);
-                                var hitsView = hitsBuffer.View.SubView(0, index);
 
-                                Span<ulong> divisorSpan = divisors.AsSpan(0, index);
+                                Span<ulong> divisorSpan = divisors.AsSpan(0, batchSize);
-                                divisorView.CopyFromCPU(ref MemoryMarshal.GetReference(divisorSpan), index);
+                                Span<byte> hitsSpan = hits.AsSpan(0, batchSize);
-                                hitsView.MemSetToZero();
 
-                                kernel(index, prime, divisorView, remainderView, hitsView);
+                                divisorView.CopyFromCPU(ref MemoryMarshal.GetReference(divisorSpan), batchSize);
+                                kernel(batchSize, prime, divisorView, hitsView);
                                 accelerator.Synchronize();
 
-                                Span<ulong> remainderSpan = remainders.AsSpan(0, index);
+                                hitsView.CopyToCPU(ref MemoryMarshal.GetReference(hitsSpan), batchSize);
-                                Span<byte> hitsSpan = hits.AsSpan(0, index);
-                                remainderView.CopyToCPU(ref MemoryMarshal.GetReference(remainderSpan), index);
-                                hitsView.CopyToCPU(ref MemoryMarshal.GetReference(hitsSpan), index);
 
-                                for (int i = 0; i < index; i++)
+                                for (int i = 0; i < batchSize; i++)
                                 {
-                                        ulong newDivisor = divisorSpan[i];
+                                        lastProcessed = divisorSpan[i];
-                                        _divisors.Add(newDivisor);
+                                        if (hitsSpan[i] != 0)
-                                        _remainders.Add(remainderSpan[i]);
-                                        _lastPrimes.Add(prime);
-                                        lastProcessed = newDivisor;
-                                        if (!composite && hitsSpan[i] != 0)
                                         {
                                                 composite = true;
+                                                break;
                                         }
                                 }
 
@@ -262,24 +157,25 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester
                                 {
                                         break;
                                 }
+
+                                if (reachedEndInBatch)
+                                {
+                                        processedAll = true;
+                                        break;
+                                }
                         }
+
+                        coveredRange = composite || processedAll || divisor > allowedMax;
                 }
                 finally
                 {
-                        ArrayPool<ulong>.Shared.Return(divisors, clearArray: true);
-                        ArrayPool<ulong>.Shared.Return(remainders, clearArray: true);
                         ArrayPool<byte>.Shared.Return(hits, clearArray: true);
+                        ArrayPool<ulong>.Shared.Return(divisors, clearArray: true);
                         hitsBuffer.Dispose();
-                        remaindersBuffer.Dispose();
                         divisorsBuffer.Dispose();
                         gpuLease.Dispose();
                 }
 
-                if (lastProcessed != 0UL)
-                {
-                        ReportStatus(lastProcessed);
-                }
-
                 return composite;
         }
 
@@ -325,7 +221,7 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester
                 return Math.Min(maxByPrime, _divisorLimit);
         }
 
-        private static void UpdateKernel(Index1D index, ulong prime, ArrayView<ulong> divisors, ArrayView<ulong> remainders, ArrayView<ulong> lastPrimes, ArrayView<byte> hits)
+        private static void CheckKernel(Index1D index, ulong prime, ArrayView<ulong> divisors, ArrayView<byte> hits)
         {
                 ulong divisor = divisors[index];
                 if (divisor <= 1UL || (divisor & 1UL) == 0UL)
@@ -334,44 +230,8 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester
                         return;
                 }
 
-                ulong lastPrime = lastPrimes[index];
-                ulong remainder = remainders[index] % divisor;
-                ulong step = prime >= lastPrime ? prime - lastPrime : 0UL;
-                if (step == 0UL)
-                {
-                        hits[index] = remainder == 0UL ? (byte)1 : (byte)0;
-                        return;
-                }
-
-                ulong pow = Pow2Mod(step, divisor);
-                ulong temp = remainder + 1UL;
-                if (temp >= divisor)
-                {
-                        temp -= divisor;
-                }
-
-                ulong product = MulMod(temp, pow, divisor);
-                ulong newRemainder = product == 0UL ? divisor - 1UL : product - 1UL;
-                bool isZero = newRemainder == 0UL;
-                remainders[index] = newRemainder;
-                lastPrimes[index] = prime;
-                hits[index] = isZero ? (byte)1 : (byte)0;
-        }
-
-        private static void InitialKernel(Index1D index, ulong prime, ArrayView<ulong> divisors, ArrayView<ulong> remainders, ArrayView<byte> hits)
-        {
-                ulong divisor = divisors[index];
-                if (divisor <= 1UL || (divisor & 1UL) == 0UL)
-                {
-                        remainders[index] = 0UL;
-                        hits[index] = 0;
-                        return;
-                }
-
                 ulong pow = Pow2Mod(prime, divisor);
-                ulong remainder = pow == 0UL ? divisor - 1UL : (pow == 1UL ? 0UL : pow - 1UL);
+                hits[index] = pow == 1UL ? (byte)1 : (byte)0;
-                remainders[index] = remainder;
-                hits[index] = remainder == 0UL ? (byte)1 : (byte)0;
         }
 
         private static ulong Pow2Mod(ulong exponent, ulong modulus)