Reintroduce MulHigh CPU helper and tighten GPU sieve loop

EvenPerfectBitScanner.Benchmarks/MulHighBenchmarks.cs

@@ -3,7 +3,6 @@ using System.Runtime.CompilerServices;
 using BenchmarkDotNet.Attributes;
 using BenchmarkDotNet.Jobs;
 using PerfectNumbers.Core;
-using PerfectNumbers.Core.Gpu;
 
 namespace EvenPerfectBitScanner.Benchmarks;
 
@@ -37,6 +36,7 @@ public class MulHighBenchmarks
     [Benchmark(Baseline = true)]
     public ulong CurrentImplementation()
     {
+        // Routes through the UInt128-based helper adopted by CPU hot paths after benchmarking.
         return ULongExtensions.MulHighCpu(Input.X, Input.Y);
     }
 
@@ -105,9 +105,7 @@ public class MulHighBenchmarks
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private static ulong MulHighWithGpuUInt128(ulong x, ulong y)
     {
-        GpuUInt128 product = new(x);
-        GpuUInt128.Mul64(ref product, 0UL, y);
-        return product.High;
+        return ULongExtensions.MulHighGpu(x, y);
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]

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<ulong>, ArrayView<byte>> PrimeTesterSmallPrimeKernel =
+    private static readonly Action<Index1D, ArrayView<ulong>, ArrayView<uint>, 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<ulong>, ArrayView<byte>>(PrimeTesterSmallPrimeKernel);
+        var kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<byte>>(PrimeTesterSmallPrimeKernel);
         _ = kernel;
     }
 

PerfectNumbers.Core.Tests/Gpu/PrimeTesterGpuTests.cs

@@ -1,6 +1,4 @@
 using FluentAssertions;
-using ILGPU;
-using ILGPU.Runtime;
 using PerfectNumbers.Core.Gpu;
 using Xunit;
 
@@ -29,34 +27,6 @@ public class PrimeTesterGpuTests
         }
     }
 
-    [Fact]
-    public void IsPrimeBatchGpu_with_preallocated_resources_marks_expected_values()
-    {
-        GpuContextPool.ForceCpu = false;
-        using var context = Context.CreateDefault();
-        using var accelerator = context.GetPreferredDevice(false).CreateAccelerator(context);
-
-        try
-        {
-            ulong[] values = [31UL, 33UL, 37UL, 39UL];
-            byte[] results = new byte[values.Length];
-
-            using var inputBuffer = accelerator.Allocate1D<ulong>(values.Length);
-            using var outputBuffer = accelerator.Allocate1D<byte>(values.Length);
-            ulong[] staging = new ulong[values.Length];
-
-            PrimeTester.IsPrimeBatchGpu(values, results, accelerator, inputBuffer, outputBuffer, staging);
-
-            results.Should().Equal(new byte[] { 1, 0, 1, 0 });
-        }
-        finally
-        {
-            PrimeTester.ClearGpuCaches(accelerator);
-            GpuContextPool.ForceCpu = false;
-            GpuContextPool.DisposeAll();
-        }
-    }
-
     [Fact]
     public void IsPrimeGpu_rejects_composites()
     {

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];

PerfectNumbers.Core/PrimeOrderCalculator.Cpu.cs

@@ -3,10 +3,9 @@ using System.Buffers;
 using System.Diagnostics;
 using System.Numerics;
 using System.Runtime.CompilerServices;
-using ILGPU;
-using ILGPU.Runtime;
-using PerfectNumbers.Core.Gpu;
 using System.Runtime.InteropServices;
+using ILGPU;
+using PerfectNumbers.Core.Gpu;
 
 namespace PerfectNumbers.Core;
 
@@ -44,90 +43,6 @@ internal static partial class PrimeOrderCalculator
 		}
 	}
 
-        private sealed class PartialFactorGpuContext : IDisposable
-        {
-                private readonly GpuKernelLease _lease;
-                private readonly GpuKernelLease.ExecutionScope _executionScope;
-                private readonly PrimeTester.KernelState _primeState;
-                private readonly PrimeTester.KernelState.ScratchBuffers _primeScratch;
-                private readonly ulong[] _primeHostBuffer;
-                private readonly byte[] _primeFlagBuffer;
-                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,
-                        GpuKernelLease.ExecutionScope executionScope,
-                        SmallPrimeFactorScratch scratch,
-                        PrimeTester.KernelState primeState,
-                        PrimeTester.KernelState.ScratchBuffers primeScratch,
-                        ulong[] primeHostBuffer,
-                        byte[] primeFlagBuffer)
-                {
-                        _lease = lease;
-                        _executionScope = executionScope;
-                        Scratch = scratch;
-                        _primeState = primeState;
-                        _primeScratch = primeScratch;
-                        _primeHostBuffer = primeHostBuffer;
-                        _primeFlagBuffer = primeFlagBuffer;
-                        _primeKernel = primeState.Kernel;
-                        _primeDevicePrimes = primeState.DevicePrimes.View;
-                        _primeDevicePrimeSquares = primeState.DevicePrimeSquares.View;
-                }
-
-                public Accelerator Accelerator => _lease.Accelerator;
-
-                public AcceleratorStream Stream => _lease.Stream;
-
-                public SmallPrimeFactorScratch Scratch { get; }
-
-                public Span<ulong> PrimeHostSpan => _primeHostBuffer.AsSpan(0, _primeScratch.Capacity);
-
-                public Span<byte> PrimeFlagSpan => _primeFlagBuffer.AsSpan(0, _primeScratch.Capacity);
-
-                public MemoryBuffer1D<ulong, Stride1D.Dense> PrimeDeviceInput => _primeScratch.Input;
-
-                public MemoryBuffer1D<byte, Stride1D.Dense> PrimeDeviceOutput => _primeScratch.Output;
-
-                public static PartialFactorGpuContext Rent()
-                {
-                        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
-                        var execution = lease.EnterExecutionScope();
-                        Accelerator accelerator = lease.Accelerator;
-                        SmallPrimeFactorScratch scratch = GpuKernelPool.EnsureSmallPrimeFactorScratch(accelerator, GpuSmallPrimeFactorSlots);
-                        PrimeTester.KernelState primeState = PrimeTester.GpuKernelState.GetOrCreate(accelerator);
-                        PrimeTester.KernelState.ScratchBuffers primeScratch = primeState.RentScratch(1, accelerator);
-                        ulong[] primeHostBuffer = ThreadStaticPools.UlongPool.Rent(primeScratch.Capacity);
-                        byte[] primeFlagBuffer = ThreadStaticPools.BytePool.Rent(primeScratch.Capacity);
-                        return new PartialFactorGpuContext(lease, execution, scratch, primeState, primeScratch, primeHostBuffer, primeFlagBuffer);
-                }
-
-                public bool CheckPrime(ulong value)
-                {
-                        Span<ulong> host = PrimeHostSpan;
-                        host[0] = value;
-                        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(), _primeDevicePrimeSquares.AsContiguous(), outputSlice.AsContiguous());
-                        Accelerator.Synchronize();
-                        Span<byte> flags = PrimeFlagSpan;
-                        outputSlice.CopyToCPU(ref flags[0], 1);
-                        return flags[0] != 0;
-                }
-
-                public void Dispose()
-                {
-                        _primeState.ReturnScratch(_primeScratch);
-                        ThreadStaticPools.UlongPool.Return(_primeHostBuffer, clearArray: false);
-                        ThreadStaticPools.BytePool.Return(_primeFlagBuffer, clearArray: false);
-                        _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.
@@ -143,20 +58,18 @@ internal static partial class PrimeOrderCalculator
 			return gpuOrder;
 		}
 
-		PartialFactorGpuContext gpuContext = PartialFactorGpuContext.Rent();
-		PartialFactorResult phiFactors = PartialFactor(phi, config, gpuContext);
+		PartialFactorResult phiFactors = PartialFactor(phi, config);
 
+		ulong result;
 		if (phiFactors.Factors is null)
 		{
-			ulong fallback = CalculateByFactorizationCpu(prime, divisorData);
+			result = CalculateByFactorizationCpu(prime, divisorData);
 			phiFactors.Dispose();
-			gpuContext.Dispose();
-			return fallback;
+			return result;
 		}
 
-		ulong result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors, gpuContext);
+		result = RunHeuristicPipelineCpu(prime, previousOrder, config, divisorData, phi, phiFactors);
 		phiFactors.Dispose();
-		gpuContext.Dispose();
 		return result;
 	}
 
@@ -166,8 +79,7 @@ internal static partial class PrimeOrderCalculator
 		in PrimeOrderSearchConfig config,
 		in MontgomeryDivisorData divisorData,
 		ulong phi,
-		PartialFactorResult phiFactors,
-		PartialFactorGpuContext gpuContext)
+		PartialFactorResult phiFactors)
 	{
 		if (phiFactors.FullyFactored && TrySpecialMaxCpu(phi, prime, phiFactors, divisorData))
 		{
@@ -177,7 +89,7 @@ internal static partial class PrimeOrderCalculator
 		ulong candidateOrder = InitializeStartingOrderCpu(prime, phi, divisorData);
 		candidateOrder = ExponentLoweringCpu(candidateOrder, prime, phiFactors, divisorData);
 
-		if (TryConfirmOrderCpu(prime, candidateOrder, divisorData, config, gpuContext, out PartialFactorResult? orderFactors))
+		if (TryConfirmOrderCpu(prime, candidateOrder, divisorData, config, out PartialFactorResult? orderFactors))
 		{
 			return candidateOrder;
 		}
@@ -188,7 +100,7 @@ internal static partial class PrimeOrderCalculator
 			return CalculateByFactorizationCpu(prime, divisorData);
 		}
 
-		if (TryHeuristicFinishCpu(prime, candidateOrder, previousOrder, divisorData, config, phiFactors, orderFactors, gpuContext, out ulong order))
+		if (TryHeuristicFinishCpu(prime, candidateOrder, previousOrder, divisorData, config, phiFactors, orderFactors, out ulong order))
 		{
 			return order;
 		}
@@ -418,7 +330,6 @@ internal static partial class PrimeOrderCalculator
 		ulong order,
 		in MontgomeryDivisorData divisorData,
 		in PrimeOrderSearchConfig config,
-		PartialFactorGpuContext gpuContext,
 		out PartialFactorResult? reusableFactorization)
 	{
 		reusableFactorization = null;
@@ -433,7 +344,7 @@ internal static partial class PrimeOrderCalculator
 			return false;
 		}
 
-		PartialFactorResult factorization = PartialFactor(order, config, gpuContext);
+		PartialFactorResult factorization = PartialFactor(order, config);
 		if (factorization.Factors is null)
 		{
 			factorization.Dispose();
@@ -581,7 +492,6 @@ internal static partial class PrimeOrderCalculator
 		in PrimeOrderSearchConfig config,
 		PartialFactorResult phiFactors,
 		PartialFactorResult? cachedOrderFactors,
-		PartialFactorGpuContext gpuContext,
 		out ulong result)
 	{
 		result = 0UL;
@@ -592,7 +502,7 @@ internal static partial class PrimeOrderCalculator
 		}
 
 		// Reuse the partial factorization from TryConfirmOrderCpu when available.
-		PartialFactorResult orderFactors = cachedOrderFactors ?? PartialFactor(order, config, gpuContext);
+		PartialFactorResult orderFactors = cachedOrderFactors ?? PartialFactor(order, config);
 		try
 		{
 			if (orderFactors.Factors is null)
@@ -708,7 +618,7 @@ internal static partial class PrimeOrderCalculator
 
 					bool equalsOne = powStepper.InitializeCpuIsUnity(candidate);
 					powStepperInitialized = true;
-					if (equalsOne && TryConfirmCandidateCpu(prime, candidate, divisorData, config, gpuContext, ref powUsed, powBudget))
+					if (equalsOne && TryConfirmCandidateCpu(prime, candidate, divisorData, config, ref powUsed, powBudget))
 					{
 						if (gpuPool is not null)
 						{
@@ -745,7 +655,7 @@ internal static partial class PrimeOrderCalculator
 						continue;
 					}
 
-					if (!TryConfirmCandidateCpu(prime, candidate, divisorData, config, gpuContext, ref powUsed, powBudget))
+					if (!TryConfirmCandidateCpu(prime, candidate, divisorData, config, ref powUsed, powBudget))
 					{
 						continue;
 					}
@@ -948,9 +858,9 @@ internal static partial class PrimeOrderCalculator
 		}
 	}
 
-	private static bool TryConfirmCandidateCpu(ulong prime, ulong candidate, in MontgomeryDivisorData divisorData, in PrimeOrderSearchConfig config, PartialFactorGpuContext gpuContext, ref int powUsed, int powBudget)
+	private static bool TryConfirmCandidateCpu(ulong prime, ulong candidate, in MontgomeryDivisorData divisorData, in PrimeOrderSearchConfig config, ref int powUsed, int powBudget)
 	{
-		PartialFactorResult factorization = PartialFactor(candidate, config, gpuContext);
+		PartialFactorResult factorization = PartialFactor(candidate, config);
 		try
 		{
 			if (factorization.Factors is null)
@@ -1114,7 +1024,7 @@ internal static partial class PrimeOrderCalculator
 	// private static ulong _partialFactorPendingHits;
 	// private static ulong _partialFactorCofactorHits;
 
-	private static PartialFactorResult PartialFactor(ulong value, in PrimeOrderSearchConfig config, PartialFactorGpuContext gpuContext)
+	private static PartialFactorResult PartialFactor(ulong value, in PrimeOrderSearchConfig config)
 	{
 		if (value <= 1UL)
 		{
@@ -1141,21 +1051,16 @@ internal static partial class PrimeOrderCalculator
 		uint limit = config.SmallFactorLimit == 0 ? uint.MaxValue : config.SmallFactorLimit;
 		ulong remaining = value;
 
-		bool gpuFactored = PrimeOrderGpuHeuristics.TryPartialFactor(
-			gpuContext.Accelerator,
-			gpuContext.Stream,
-			gpuContext.Scratch,
-			value,
-			limit,
-			primeSlots,
-			exponentSlots,
-			out factorCount,
-			out remaining,
-			out _);
+		bool gpuFactored = false;
+		if (IsGpuHeuristicDevice)
+		{
+			gpuFactored = PrimeOrderGpuHeuristics.TryPartialFactor(value, limit, primeSlots, exponentSlots, out factorCount, out remaining, out _);
+		}
 
 		List<PendingEntry> pending = ThreadStaticPools.RentPrimeOrderPendingEntryList(2);
 		Stack<ulong>? compositeStack = null;
 		PartialFactorResult result;
+		PrimeTester primeTester = PrimeTester.Exclusive;
 
 		if (!gpuFactored)
 		{
@@ -1228,7 +1133,7 @@ internal static partial class PrimeOrderCalculator
 					// Atomic.Add(ref _partialFactorHits, 1UL);
 					// Console.WriteLine($"Partial factor hits {Volatile.Read(ref _partialFactorHits)}");
 
-					bool isPrime = gpuContext.CheckPrime(composite);
+					bool isPrime = PrimeTester.IsPrimeGpu(composite);
 					// bool isPrime = Open.Numeric.Primes.Prime.Numbers.IsPrime(composite);
 
 					if (isPrime)
@@ -1295,7 +1200,7 @@ internal static partial class PrimeOrderCalculator
 				// Atomic.Add(ref _partialFactorPendingHits, 1UL);
 				// Console.WriteLine($"Partial factor pending hits {Volatile.Read(ref _partialFactorPendingHits)}");
 
-				bool isPrime = gpuContext.CheckPrime(composite);
+				bool isPrime = PrimeTester.IsPrimeGpu(composite);
 				// bool isPrime = Open.Numeric.Primes.Prime.Numbers.IsPrime(composite);
 
 				entry = entry.WithPrimality(isPrime);
@@ -1332,7 +1237,7 @@ internal static partial class PrimeOrderCalculator
 			// Atomic.Add(ref _partialFactorCofactorHits, 1UL);
 			// Console.WriteLine($"Partial factor cofactor hits {Volatile.Read(ref _partialFactorCofactorHits)}");
 
-			cofactorIsPrime = gpuContext.CheckPrime(cofactor);
+			cofactorIsPrime = PrimeTester.IsPrimeCpu(cofactor, CancellationToken.None);
 			// cofactorIsPrime = Open.Numeric.Primes.Prime.Numbers.IsPrime(cofactor);
 		}
 
@@ -1425,6 +1330,7 @@ internal static partial class PrimeOrderCalculator
 		return result;
 	}
 
+
 	private static bool TryPopulateSmallPrimeFactorsCpu(
 			ulong value,
 			uint limit,

PerfectNumbers.Core/PrimeOrderCalculator.Gpu.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Runtime.InteropServices;
 using ILGPU;
 using ILGPU.Runtime;
 using ILGPU.Runtime.OpenCL;
@@ -26,45 +27,63 @@ internal static partial class PrimeOrderCalculator
 		{
 			var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
 			var execution = lease.EnterExecutionScope();
-			Accelerator accelerator = lease.Accelerator;
-			AcceleratorStream stream = lease.Stream;
-
-			SmallPrimeFactorScratch scratch = GpuKernelPool.EnsureSmallPrimeFactorScratch(accelerator, GpuSmallPrimeFactorSlots);
-
-			bool success = PrimeOrderGpuHeuristics.TryPartialFactor(
-				accelerator,
-				stream,
-				scratch,
-				value,
-				limit,
-				primeBuffer,
-				exponentBuffer,
-				out factorCount,
-				out remaining,
-				out _);
-
-			execution.Dispose();
-			lease.Dispose();
-
-			if (!success)
+			try
 			{
-				return false;
-			}
+				Accelerator accelerator = lease.Accelerator;
+				AcceleratorStream stream = lease.Stream;
 
-			for (int i = 0; i < factorCount; i++)
+				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);
+
+				stream.Synchronize();
+
+				factorCount = 0;
+				scratch.CountSlot.View.CopyToCPU(ref factorCount, 1);
+				factorCount = Math.Min(factorCount, GpuSmallPrimeFactorSlots);
+
+				if (factorCount > 0)
+				{
+					scratch.PrimeSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(primeBuffer), factorCount);
+					scratch.ExponentSlots.View.CopyToCPU(ref MemoryMarshal.GetReference(exponentBuffer), factorCount);
+				}
+
+				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;
+			}
+			finally
 			{
-				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);
+				execution.Dispose();
+				lease.Dispose();
 			}
-
-			return true;
 		}
 		finally
 		{

PerfectNumbers.Core/PrimeTester.cs

@@ -3,7 +3,6 @@ using System.Buffers;
 using System.Collections.Concurrent;
 using System.Numerics;
 using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
 using System.Threading;
 using ILGPU;
 using ILGPU.Algorithms;
@@ -113,6 +112,7 @@ public sealed class PrimeTester
         var gpu = PrimeTesterGpuContextPool.Rent();
         var accelerator = gpu.Accelerator;
         var state = GpuKernelState.GetOrCreate(accelerator);
+        int totalLength = values.Length;
         int batchSize = Math.Max(1, GpuBatchSize);
 
         lock (gpu.ExecutionLock)
@@ -123,67 +123,30 @@ public sealed class PrimeTester
             ArrayPool<ulong> pool = ThreadStaticPools.UlongPool;
             ulong[] temp = pool.Rent(batchSize);
 
-            try
+            int pos = 0;
+            while (pos < totalLength)
             {
-                IsPrimeBatchGpu(values, results, accelerator, input, output, temp);
-            }
-            finally
-            {
-                pool.Return(temp, clearArray: false);
-                state.ReturnScratch(scratch);
+                int remaining = totalLength - pos;
+                int count = remaining > batchSize ? batchSize : remaining;
+
+                values.Slice(pos, count).CopyTo(temp);
+                input.View.CopyFromCPU(ref temp[0], count);
+
+                state.Kernel(count, input.View, state.DevicePrimes.View, output.View);
+                accelerator.Synchronize();
+                output.View.CopyToCPU(ref results[pos], count);
+
+                pos += count;
             }
+
+            pool.Return(temp, clearArray: false);
+            state.ReturnScratch(scratch);
         }
 
         gpu.Dispose();
         limiter.Dispose();
     }
 
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    public static void IsPrimeBatchGpu(
-        ReadOnlySpan<ulong> values,
-        Span<byte> results,
-        Accelerator accelerator,
-        MemoryBuffer1D<ulong, Stride1D.Dense> deviceInput,
-        MemoryBuffer1D<byte, Stride1D.Dense> deviceOutput,
-        Span<ulong> stagingBuffer)
-    {
-        if (values.IsEmpty)
-        {
-            return;
-        }
-
-        int capacity = (int)Math.Min(Math.Min(deviceInput.Length, deviceOutput.Length), stagingBuffer.Length);
-
-        var state = GpuKernelState.GetOrCreate(accelerator);
-        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);
-
-        int totalLength = values.Length;
-        int pos = 0;
-        while (pos < totalLength)
-        {
-            int remaining = totalLength - pos;
-            int count = remaining > capacity ? capacity : remaining;
-
-            var stagingSlice = stagingSpan.Slice(0, count);
-            values.Slice(pos, count).CopyTo(stagingSlice);
-
-            var inputSlice = inputView.SubView(0, count);
-            inputSlice.CopyFromCPU(ref stagingRef, count);
-
-            var outputSlice = outputView.SubView(0, count);
-            state.Kernel(count, inputSlice, primesView, primeSquaresView, outputSlice);
-            accelerator.Synchronize();
-            outputSlice.CopyToCPU(ref results[pos], count);
-
-            pos += count;
-        }
-    }
-
     internal static class PrimeTesterGpuContextPool
     {
         internal sealed class PooledContext : IDisposable
@@ -273,10 +236,9 @@ 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<ulong>, ArrayView<byte>> Kernel { get; }
+        public Action<Index1D, ArrayView<ulong>, ArrayView<uint>, 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
@@ -287,15 +249,12 @@ public sealed class PrimeTester
         {
             _accel = accelerator;
             // Compile once per accelerator and upload primes once.
-            Kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, ArrayView<ulong>, ArrayView<byte>>(PrimeTesterKernels.SmallPrimeSieveKernel);
+            Kernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView<ulong>, ArrayView<uint>, 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
@@ -349,7 +308,6 @@ public sealed class PrimeTester
             }
 
             DevicePrimes.Dispose();
-            DevicePrimeSquares.Dispose();
         }
     }
 

PerfectNumbers.Core/ULongExtensions.cs

@@ -441,8 +441,8 @@ public static partial class ULongExtensions
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]
 	public static ulong MulHighGpu(this ulong x, ulong y)
 	{
-		// The CPU helpers rely on UInt128 via MulHighCpu; retain this manual decomposition for GPU-style
-		// arithmetic so callers can avoid the intrinsic when staging work for accelerators.
+		// Retain this manual decomposition for GPU-style arithmetic so callers staging accelerator work
+		// can avoid the UInt128 intrinsic while keeping parity with GpuUInt128.MulHigh.
 		ulong xLow = (uint)x;
 		ulong xHigh = x >> 32;
 		ulong yLow = (uint)y;