Merge branch 'codex/optimize-evenperfectbitscanner-for-static-arrays' from github

PerfectNumbers.Core/Gpu/GpuContextPool.cs

@@ -1,4 +1,3 @@
-using System;
 using System.Collections.Concurrent;
 using ILGPU;
 using ILGPU.Runtime;
@@ -15,14 +14,11 @@ public static class GpuContextPool
 		public readonly Context Context;
 		public readonly Accelerator Accelerator;
 
-		private static readonly Context _sharedContext = Context.CreateDefault();
-		private static readonly Device _sharedDevice = _sharedContext.GetPreferredDevice(false);
-		private static readonly Accelerator _sharedAccelerator = _sharedDevice.CreateAccelerator(_sharedContext);
-
 		public PooledContext()
 		{
-			Context = _sharedContext;
-			Accelerator = _sharedAccelerator;
+			Context = SharedGpuContext.Context;
+			Accelerator = SharedGpuContext.Accelerator;
+			GpuStaticTableInitializer.EnsureStaticTables(Accelerator);
 			// NOTE: Avoid loading/compiling any kernel here to prevent implicit
 			// CL stream/queue creation during accelerator construction.
 			// Some OpenCL drivers are fragile when a queue is created immediately
@@ -48,7 +44,7 @@ public static class GpuContextPool
 	private static int WarmedGpuContextCount;
 
 	// Allows callers to choose CPU/GPU per use-case, decoupled from ForceCpu.
-        	public static GpuContextLease Rent()
+	public static GpuContextLease Rent()
 	{
 		if (GpuPool.TryDequeue(out var gpu))
 		{

PerfectNumbers.Core/Gpu/GpuKernelLease.cs

@@ -2,6 +2,7 @@ using System;
 using System.Collections.Concurrent;
 using ILGPU;
 using ILGPU.Runtime;
+using PerfectNumbers.Core;
 using static PerfectNumbers.Core.Gpu.GpuContextPool;
 
 namespace PerfectNumbers.Core.Gpu;
@@ -121,6 +122,20 @@ public sealed class GpuKernelLease
         }
     }
 
+    public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>> SpecialMaxKernel
+    {
+        get
+        {
+            var accel = Accelerator;
+            return KernelContainer.InitOnce(ref _kernels!.SpecialMax, () =>
+            {
+                var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>(PrimeOrderGpuHeuristics.EvaluateSpecialMaxCandidatesKernel);
+                var kernel = KernelUtil.GetKernel(loaded);
+                return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>>();
+            });
+        }
+    }
+
     public Action<AcceleratorStream, Index1D, ulong, uint, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>> SmallPrimeFactorKernel
     {
         get

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -1,3 +1,4 @@
+using PerfectNumbers.Core;
 using System.Collections.Concurrent;
 using ILGPU;
 using ILGPU.Runtime;
@@ -37,6 +38,7 @@ public sealed class KernelContainer
     public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
         ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>? Pow2ModOrder;
     public Action<AcceleratorStream, Index1D, ulong, uint, ArrayView1D<uint, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>>? SmallPrimeFactor;
+    public Action<AcceleratorStream, Index1D, ulong, ArrayView1D<ulong, Stride1D.Dense>, int, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>, ArrayView1D<byte, Stride1D.Dense>>? SpecialMax;
 
     // Optional device buffer with small divisor cycles (<= 4M). Index = divisor, value = cycle length.
     public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallCycles;
@@ -46,6 +48,9 @@ public sealed class KernelContainer
     public MemoryBuffer1D<int, Stride1D.Dense>? SmallPrimeFactorExponentSlots;
     public MemoryBuffer1D<int, Stride1D.Dense>? SmallPrimeFactorCountSlot;
     public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimeFactorRemainingSlot;
+    public MemoryBuffer1D<ulong, Stride1D.Dense>? SpecialMaxFactors;
+    public MemoryBuffer1D<ulong, Stride1D.Dense>? SpecialMaxCandidates;
+    public MemoryBuffer1D<byte, Stride1D.Dense>? SpecialMaxResult;
     public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastOne;
     public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallPrimesPow2LastOne;
     public MemoryBuffer1D<uint, Stride1D.Dense>? SmallPrimesLastSeven;
@@ -129,6 +134,29 @@ public readonly struct SmallPrimeFactorScratch
     }
 }
 
+public readonly struct SpecialMaxScratch
+{
+    public readonly MemoryBuffer1D<ulong, Stride1D.Dense> FactorValues;
+    public readonly MemoryBuffer1D<ulong, Stride1D.Dense> CandidateValues;
+    public readonly MemoryBuffer1D<byte, Stride1D.Dense> ResultSlot;
+
+    public SpecialMaxScratch(
+        MemoryBuffer1D<ulong, Stride1D.Dense> factorValues,
+        MemoryBuffer1D<ulong, Stride1D.Dense> candidateValues,
+        MemoryBuffer1D<byte, Stride1D.Dense> resultSlot)
+    {
+        FactorValues = factorValues;
+        CandidateValues = candidateValues;
+        ResultSlot = resultSlot;
+    }
+
+    public ArrayView1D<ulong, Stride1D.Dense> FactorsView => FactorValues.View;
+
+    public ArrayView1D<ulong, Stride1D.Dense> CandidatesView => CandidateValues.View;
+
+    public ArrayView1D<byte, Stride1D.Dense> ResultView => ResultSlot.View;
+}
+
 public class GpuKernelPool
 {
     private static readonly ConcurrentDictionary<Accelerator, KernelContainer> KernelCache = new(); // TODO: Replace this concurrent map with a simple accelerator-indexed lookup once kernel launchers are prewarmed during startup so we can drop the thread-safe wrapper entirely.
@@ -138,6 +166,13 @@ public class GpuKernelPool
         return KernelCache.GetOrAdd(accelerator, _ => new KernelContainer());
     }
 
+    internal static void PreloadStaticTables(Accelerator accelerator)
+    {
+        EnsureSmallCyclesOnDevice(accelerator);
+        _ = EnsureSmallPrimesOnDevice(accelerator);
+        _ = EnsureSmallPrimeFactorTables(accelerator);
+    }
+
     // Ensures the small cycles table is uploaded to the device for the given accelerator.
     // Returns the ArrayView to pass into kernels (when kernels are extended to accept it).
     public static ArrayView1D<ulong, Stride1D.Dense> EnsureSmallCyclesOnDevice(Accelerator accelerator)
@@ -277,6 +312,43 @@ public class GpuKernelPool
         }
     }
 
+    public static SpecialMaxScratch EnsureSpecialMaxScratch(Accelerator accelerator, int factorCapacity)
+    {
+        if (factorCapacity <= 0)
+        {
+            factorCapacity = 1;
+        }
+
+        var kernels = GetKernels(accelerator);
+        lock (kernels)
+        {
+            MemoryBuffer1D<ulong, Stride1D.Dense>? factorValues = kernels.SpecialMaxFactors;
+            if (factorValues is null || factorValues.Length < factorCapacity)
+            {
+                factorValues?.Dispose();
+                factorValues = accelerator.Allocate1D<ulong>(factorCapacity);
+                kernels.SpecialMaxFactors = factorValues;
+            }
+
+            MemoryBuffer1D<ulong, Stride1D.Dense>? candidateValues = kernels.SpecialMaxCandidates;
+            if (candidateValues is null || candidateValues.Length < factorCapacity)
+            {
+                candidateValues?.Dispose();
+                candidateValues = accelerator.Allocate1D<ulong>(factorCapacity);
+                kernels.SpecialMaxCandidates = candidateValues;
+            }
+
+            MemoryBuffer1D<byte, Stride1D.Dense>? resultSlot = kernels.SpecialMaxResult;
+            if (resultSlot is null)
+            {
+                resultSlot = accelerator.Allocate1D<byte>(1);
+                kernels.SpecialMaxResult = resultSlot;
+            }
+
+            return new SpecialMaxScratch(factorValues!, candidateValues!, resultSlot!);
+        }
+    }
+
     public static GpuKernelLease GetKernel(bool useGpuOrder)
     {
         var limiter = GpuPrimeWorkLimiter.Acquire();

PerfectNumbers.Core/Gpu/GpuStaticTableInitializer.cs

@@ -0,0 +1,44 @@
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Runtime.CompilerServices;
+using ILGPU;
+using ILGPU.Runtime;
+using PerfectNumbers.Core;
+
+namespace PerfectNumbers.Core.Gpu;
+
+internal static class GpuStaticTableInitializer
+{
+    private static readonly ConcurrentDictionary<Accelerator, byte> Initialized = new(AcceleratorReferenceComparer.Instance);
+
+    internal static void EnsureStaticTables(Accelerator accelerator)
+    {
+        if (accelerator is null)
+        {
+            return;
+        }
+
+        Initialized.GetOrAdd(accelerator, static acc =>
+        {
+            PrimeTester.PrimeTesterGpuContextPool.EnsureStaticTables(acc);
+            GpuKernelPool.PreloadStaticTables(acc);
+            PrimeOrderGpuHeuristics.PreloadStaticTables(acc);
+            return 0;
+        });
+    }
+
+    private sealed class AcceleratorReferenceComparer : IEqualityComparer<Accelerator>
+    {
+        internal static AcceleratorReferenceComparer Instance { get; } = new();
+
+        public bool Equals(Accelerator? x, Accelerator? y)
+        {
+            return ReferenceEquals(x, y);
+        }
+
+        public int GetHashCode(Accelerator obj)
+        {
+            return RuntimeHelpers.GetHashCode(obj);
+        }
+    }
+}

PerfectNumbers.Core/Gpu/Kernels/PrimeOrderSpecialMaxKernels.cs

@@ -0,0 +1,80 @@
+using ILGPU;
+using ILGPU.Runtime;
+using PerfectNumbers.Core;
+
+namespace PerfectNumbers.Core.Gpu;
+
+internal static partial class PrimeOrderGpuHeuristics
+{
+    internal static void EvaluateSpecialMaxCandidatesKernel(
+        Index1D index,
+        ulong phi,
+        ArrayView1D<ulong, Stride1D.Dense> factors,
+        int factorCount,
+        MontgomeryDivisorData divisor,
+        ArrayView1D<ulong, Stride1D.Dense> candidates,
+        ArrayView1D<byte, Stride1D.Dense> resultOut)
+    {
+        if (index != 0)
+        {
+            return;
+        }
+
+        int actual = 0;
+        for (int i = 0; i < factorCount; i++)
+        {
+            ulong factor = factors[i];
+            if (factor <= 1UL)
+            {
+                continue;
+            }
+
+            ulong reduced = phi / factor;
+            if (reduced == 0UL)
+            {
+                continue;
+            }
+
+            candidates[actual] = reduced;
+            actual++;
+        }
+
+        if (actual == 0)
+        {
+            resultOut[0] = 1;
+            return;
+        }
+
+        for (int i = 1; i < actual; i++)
+        {
+            ulong current = candidates[i];
+            int insert = i - 1;
+            while (insert >= 0 && candidates[insert] > current)
+            {
+                candidates[insert + 1] = candidates[insert];
+                insert--;
+            }
+
+            candidates[insert + 1] = current;
+        }
+
+        var divisorPartial = new GpuDivisorPartialData(divisor.Modulus);
+        var stepper = new ExponentRemainderStepperGpu(divisorPartial);
+        if (stepper.InitializeIsUnityGpu(candidates[0]))
+        {
+            resultOut[0] = 0;
+            return;
+        }
+
+        for (int i = 1; i < actual; i++)
+        {
+            if (stepper.ComputeNextIsUnityGpu(candidates[i]))
+            {
+                resultOut[0] = 0;
+                return;
+            }
+        }
+
+        resultOut[0] = 1;
+    }
+}

PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs

@@ -18,6 +18,11 @@ internal enum GpuPow2ModStatus
 
 internal static partial class PrimeOrderGpuHeuristics
 {
+    internal static void PreloadStaticTables(Accelerator accelerator)
+    {
+        _ = GetSmallPrimeDeviceCache(accelerator);
+    }
+
     private static readonly ConcurrentDictionary<ulong, byte> OverflowedPrimes = new();
     private static readonly ConcurrentDictionary<UInt128, byte> OverflowedPrimesWide = new();
     private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>> Pow2ModKernelCache = new();

PerfectNumbers.Core/Gpu/SharedGpuContext.cs

@@ -0,0 +1,19 @@
+using System;
+using System.Threading;
+using ILGPU;
+using ILGPU.Runtime;
+
+namespace PerfectNumbers.Core.Gpu;
+
+internal static class SharedGpuContext
+{
+    private static readonly Lazy<Context> s_context = new(() => ILGPU.Context.CreateDefault(), LazyThreadSafetyMode.ExecutionAndPublication);
+    private static readonly Lazy<Device> s_device = new(() => s_context.Value.GetPreferredDevice(false), LazyThreadSafetyMode.ExecutionAndPublication);
+    private static readonly Lazy<Accelerator> s_accelerator = new(() => s_device.Value.CreateAccelerator(s_context.Value), LazyThreadSafetyMode.ExecutionAndPublication);
+
+    internal static Context Context => s_context.Value;
+
+    internal static Device Device => s_device.Value;
+
+    internal static Accelerator Accelerator => s_accelerator.Value;
+}