Use precomputed small prime squares in GPU sieve

PerfectNumbers.Core.Tests/Gpu/GpuKernelCompilationTests.cs

@@ -121,7 +121,7 @@ public class GpuKernelCompilationTests
     private static readonly Action<Index1D, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> MersenneDivisorCyclesKernel =
         DivisorCycleKernels.GpuDivisorCycleKernel;
 
-    private static readonly Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<byte>> PrimeTesterSmallPrimeKernel =
+    private static readonly Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<ulong>, ArrayView<byte>> PrimeTesterSmallPrimeKernel =
         PrimeTesterKernels.SmallPrimeSieveKernel;
 
     private static readonly Action<Index1D, ArrayView<ulong>, ArrayView<byte>> PrimeTesterSharesFactorKernel =
@@ -336,7 +336,7 @@ public class GpuKernelCompilationTests
 
     private static void CompilePrimeTesterSmallPrimeKernel(Accelerator accelerator)
     {
-        var kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<byte>>(PrimeTesterSmallPrimeKernel);
+        var kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterSmallPrimeKernel);
         _ = kernel;
     }
 

PerfectNumbers.Core/Gpu/Kernels/PrimeTesterKernels.cs

@@ -7,7 +7,12 @@ namespace PerfectNumbers.Core.Gpu;
 internal static class PrimeTesterKernels
 {
     // GPU kernel: small-prime sieve only. Returns 1 if passes sieve (probable prime), 0 otherwise.
-    public static void SmallPrimeSieveKernel(Index1D index, ArrayView<ulong> numbers, ArrayView<uint> smallPrimes, ArrayView<byte> results)
+    public static void SmallPrimeSieveKernel(
+        Index1D index,
+        ArrayView<ulong> numbers,
+        ArrayView<uint> smallPrimes,
+        ArrayView<ulong> smallPrimeSquares,
+        ArrayView<byte> results)
     {
         ulong n = numbers[index];
 
@@ -28,7 +33,7 @@ internal static class PrimeTesterKernels
         for (int i = 0; i < length; i++)
         {
             ulong prime = smallPrimes[i];
-            ulong primeSquare = (ulong)prime * prime;
+            ulong primeSquare = smallPrimeSquares[i];
             bool withinRange = primeSquare <= n;
             ulong remainder = n % prime;
             bool divides = remainder == 0UL;

PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs

@@ -52,8 +52,9 @@ internal static partial class PrimeOrderCalculator
                 private readonly PrimeTester.KernelState.ScratchBuffers _primeScratch;
                 private readonly ulong[] _primeHostBuffer;
                 private readonly byte[] _primeFlagBuffer;
-                private readonly Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<byte>> _primeKernel;
+                private readonly Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<ulong>, ArrayView<byte>> _primeKernel;
                 private readonly ArrayView1D<uint, Stride1D.Dense> _primeDevicePrimes;
+                private readonly ArrayView1D<ulong, Stride1D.Dense> _primeDevicePrimeSquares;
 
                 private PartialFactorGpuContext(
                         GpuKernelLease lease,
@@ -73,6 +74,7 @@ internal static partial class PrimeOrderCalculator
                         _primeFlagBuffer = primeFlagBuffer;
                         _primeKernel = primeState.Kernel;
                         _primeDevicePrimes = primeState.DevicePrimes.View;
+                        _primeDevicePrimeSquares = primeState.DevicePrimeSquares.View;
                 }
 
                 public Accelerator Accelerator => _lease.Accelerator;
@@ -109,7 +111,7 @@ internal static partial class PrimeOrderCalculator
                         ArrayView1D<ulong, Stride1D.Dense> inputSlice = PrimeDeviceInput.View.SubView(0, 1);
                         inputSlice.CopyFromCPU(ref host[0], 1);
                         ArrayView1D<byte, Stride1D.Dense> outputSlice = PrimeDeviceOutput.View.SubView(0, 1);
-                        _primeKernel(1, inputSlice.AsContiguous(), _primeDevicePrimes.AsContiguous(), outputSlice.AsContiguous());
+                        _primeKernel(1, inputSlice.AsContiguous(), _primeDevicePrimes.AsContiguous(), _primeDevicePrimeSquares.AsContiguous(), outputSlice.AsContiguous());
                         Accelerator.Synchronize();
                         Span<byte> flags = PrimeFlagSpan;
                         outputSlice.CopyToCPU(ref flags[0], 1);

PerfectNumbers.Core/PrimeTester.cs

@@ -158,6 +158,7 @@ public sealed class PrimeTester
         var inputView = deviceInput.View;
         var outputView = deviceOutput.View;
         var primesView = state.DevicePrimes.View;
+        var primeSquaresView = state.DevicePrimeSquares.View;
         var stagingSpan = stagingBuffer;
         ref ulong stagingRef = ref MemoryMarshal.GetReference(stagingSpan);
 
@@ -175,7 +176,7 @@ public sealed class PrimeTester
             inputSlice.CopyFromCPU(ref stagingRef, count);
 
             var outputSlice = outputView.SubView(0, count);
-            state.Kernel(count, inputSlice, primesView, outputSlice);
+            state.Kernel(count, inputSlice, primesView, primeSquaresView, outputSlice);
             accelerator.Synchronize();
             outputSlice.CopyToCPU(ref results[pos], count);
 
@@ -272,9 +273,10 @@ public sealed class PrimeTester
     // Per-accelerator GPU state for prime sieve (kernel + uploaded primes).
     internal sealed class KernelState
     {
-        public Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<byte>> Kernel { get; }
+        public Action<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<ulong>, ArrayView<byte>> Kernel { get; }
         public Action<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>> HeuristicTrialDivisionKernel { get; }
         public MemoryBuffer1D<uint, Stride1D.Dense> DevicePrimes { get; }
+        public MemoryBuffer1D<ulong, Stride1D.Dense> DevicePrimeSquares { get; }
         private readonly Accelerator _accel;
         private readonly System.Collections.Concurrent.ConcurrentBag<ScratchBuffers> _scratchPool = [];
         // TODO: Replace this ConcurrentBag with the lock-free ring buffer variant validated in
@@ -285,12 +287,15 @@ public sealed class PrimeTester
         {
             _accel = accelerator;
             // Compile once per accelerator and upload primes once.
-            Kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
+            Kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
             HeuristicTrialDivisionKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ulong, ArrayView<byte>>(PrimeTesterKernels.HeuristicTrialDivisionKernel);
 
             var primes = PrimesGenerator.SmallPrimes;
+            var primeSquares = PrimesGenerator.SmallPrimesPow2;
             DevicePrimes = accelerator.Allocate1D<uint>(primes.Length);
+            DevicePrimeSquares = accelerator.Allocate1D<ulong>(primeSquares.Length);
             DevicePrimes.View.CopyFromCPU(primes);
+            DevicePrimeSquares.View.CopyFromCPU(primeSquares);
         }
 
         internal sealed class ScratchBuffers : IDisposable
@@ -344,6 +349,7 @@ public sealed class PrimeTester
             }
 
             DevicePrimes.Dispose();
+            DevicePrimeSquares.Dispose();
         }
     }