Reintroduce MulHigh CPU helper and tighten GPU sieve loop

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -46,7 +46,6 @@ 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;
@@ -108,20 +107,17 @@ 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<byte, Stride1D.Dense> fullyFactoredSlot)
+        MemoryBuffer1D<ulong, Stride1D.Dense> remainingSlot)
     {
         PrimeSlots = primeSlots;
         ExponentSlots = exponentSlots;
         CountSlot = countSlot;
         RemainingSlot = remainingSlot;
-        FullyFactoredSlot = fullyFactoredSlot;
     }
 
     public void Clear()
@@ -130,7 +126,6 @@ public readonly struct SmallPrimeFactorScratch
         ExponentSlots.MemSetToZero();
         CountSlot.MemSetToZero();
         RemainingSlot.MemSetToZero();
-        FullyFactoredSlot.MemSetToZero();
     }
 }
 
@@ -278,14 +273,7 @@ public class GpuKernelPool
                 kernels.SmallPrimeFactorRemainingSlot = 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!);
+            return new SmallPrimeFactorScratch(primeSlots!, exponentSlots!, countSlot!, remainingSlot!);
         }
     }
 

PerfectNumbers.Core/Gpu/Kernels/PrimeTesterKernels.cs

@@ -7,12 +7,7 @@ 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<ulong> smallPrimeSquares,
-        ArrayView<byte> results)
+    public static void SmallPrimeSieveKernel(Index1D index, ArrayView<ulong> numbers, ArrayView<uint> smallPrimes, ArrayView<byte> results)
     {
         ulong n = numbers[index];
         byte result = 1;
@@ -27,7 +22,7 @@ internal static class PrimeTesterKernels
         for (int i = 0; i < length; i++)
         {
             ulong prime = smallPrimes[i];
-            ulong primeSquare = smallPrimeSquares[i];
+            ulong primeSquare = (ulong)prime * prime;
             if (primeSquare > n)
             {
                 break;

PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs

@@ -130,11 +130,7 @@ internal static partial class PrimeOrderGpuHeuristics
         });
     }
 
-
     public static bool TryPartialFactor(
-        Accelerator accelerator,
-        AcceleratorStream stream,
-        SmallPrimeFactorScratch scratch,
         ulong value,
         uint limit,
         Span<ulong> primeTargets,
@@ -147,63 +143,28 @@ internal static partial class PrimeOrderGpuHeuristics
         remaining = value;
         fullyFactored = false;
 
+        if (primeTargets.Length == 0 || exponentTargets.Length == 0)
+        {
+            return false;
+        }
+
         primeTargets.Clear();
         exponentTargets.Clear();
 
-        return TryPartialFactorCore(
-            accelerator,
-            stream,
-            scratch,
-            value,
-            limit,
-            primeTargets,
-            exponentTargets,
-            out factorCount,
-            out remaining,
-            out fullyFactored);
-    }
+        if (!TryLaunchPartialFactorKernel(
+                value,
+                limit,
+                primeTargets,
+                exponentTargets,
+                out int extracted,
+                out ulong leftover,
+                out bool kernelFullyFactored))
+        {
+            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)
+        if (extracted > capacity)
         {
             factorCount = 0;
             remaining = value;
@@ -211,21 +172,104 @@ internal static partial class PrimeOrderGpuHeuristics
             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;
+        factorCount = extracted;
+        remaining = leftover;
+        fullyFactored = kernelFullyFactored && leftover == 1UL;
         return true;
     }
 
+    private static bool TryLaunchPartialFactorKernel(
+        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;
+
+        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(
+                stream,
+                1,
+                cache.Primes!.View,
+                cache.Squares!.View,
+                cache.Count,
+                primeTargets.Length,
+                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;
+        }
+        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;
+        }
+    }
+
     public static GpuPow2ModStatus TryPow2Mod(ulong exponent, ulong prime, out ulong remainder, in MontgomeryDivisorData divisorData)
     {
         Span<ulong> exponents = stackalloc ulong[1];