Preload divisor generator tables for GPU pool contexts

2025-10-31 23:01:28 +01:00
1 parent 469c8057d3
commit b7c34f4960
@@ -0,0 +1,142 @@
+using System;
+using System.Collections.Concurrent;
+using System.Runtime.CompilerServices;
+using ILGPU;
+using ILGPU.Runtime;
+
+namespace PerfectNumbers.Core.Gpu;
+
+internal static class DivisorGeneratorGpuCache
+{
+        private sealed class AcceleratorReferenceComparer : IEqualityComparer<Accelerator>
+        {
+                public 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);
+                }
+        }
+
+        private sealed class DivisorGeneratorGpuTables : IDisposable
+        {
+                internal DivisorGeneratorGpuTables(Accelerator accelerator)
+                {
+                        SmallPrimes = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimes);
+                        SmallPrimesPow2 = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2);
+                        SmallPrimesLastOne = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimesLastOne);
+                        SmallPrimesPow2LastOne = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2LastOne);
+                        SmallPrimesLastSeven = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimesLastSeven);
+                        SmallPrimesPow2LastSeven = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2LastSeven);
+                        SmallPrimesLastThree = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimesLastThree);
+                        SmallPrimesPow2LastThree = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2LastThree);
+                        SmallPrimesLastNine = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimesLastNine);
+                        SmallPrimesPow2LastNine = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2LastNine);
+                        SmallPrimesLastOneWithoutLastThree = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimesLastOneWithoutLastThree);
+                        SmallPrimesPow2LastOneWithoutLastThree = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2LastOneWithoutLastThree);
+                        SmallPrimesLastSevenWithoutLastThree = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimesLastSevenWithoutLastThree);
+                        SmallPrimesPow2LastSevenWithoutLastThree = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2LastSevenWithoutLastThree);
+                        SmallPrimesLastNineWithoutLastThree = AllocateAndCopyUInt(accelerator, DivisorGenerator.SmallPrimesLastNineWithoutLastThree);
+                        SmallPrimesPow2LastNineWithoutLastThree = AllocateAndCopyULong(accelerator, DivisorGenerator.SmallPrimesPow2LastNineWithoutLastThree);
+                }
+
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimes { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2 { get; }
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimesLastOne { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2LastOne { get; }
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimesLastSeven { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2LastSeven { get; }
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimesLastThree { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2LastThree { get; }
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimesLastNine { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2LastNine { get; }
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimesLastOneWithoutLastThree { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2LastOneWithoutLastThree { get; }
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimesLastSevenWithoutLastThree { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2LastSevenWithoutLastThree { get; }
+                internal MemoryBuffer1D<uint, Stride1D.Dense> SmallPrimesLastNineWithoutLastThree { get; }
+                internal MemoryBuffer1D<ulong, Stride1D.Dense> SmallPrimesPow2LastNineWithoutLastThree { get; }
+
+                public void Dispose()
+                {
+                        SmallPrimes.Dispose();
+                        SmallPrimesPow2.Dispose();
+                        SmallPrimesLastOne.Dispose();
+                        SmallPrimesPow2LastOne.Dispose();
+                        SmallPrimesLastSeven.Dispose();
+                        SmallPrimesPow2LastSeven.Dispose();
+                        SmallPrimesLastThree.Dispose();
+                        SmallPrimesPow2LastThree.Dispose();
+                        SmallPrimesLastNine.Dispose();
+                        SmallPrimesPow2LastNine.Dispose();
+                        SmallPrimesLastOneWithoutLastThree.Dispose();
+                        SmallPrimesPow2LastOneWithoutLastThree.Dispose();
+                        SmallPrimesLastSevenWithoutLastThree.Dispose();
+                        SmallPrimesPow2LastSevenWithoutLastThree.Dispose();
+                        SmallPrimesLastNineWithoutLastThree.Dispose();
+                        SmallPrimesPow2LastNineWithoutLastThree.Dispose();
+                }
+
+                private static MemoryBuffer1D<uint, Stride1D.Dense> AllocateAndCopyUInt(Accelerator accelerator, uint[] source)
+                {
+                        var buffer = accelerator.Allocate1D<uint>(source.Length);
+                        if (source.Length > 0)
+                        {
+                                buffer.View.CopyFromCPU(source);
+                        }
+
+                        return buffer;
+                }
+
+                private static MemoryBuffer1D<ulong, Stride1D.Dense> AllocateAndCopyULong(Accelerator accelerator, ulong[] source)
+                {
+                        var buffer = accelerator.Allocate1D<ulong>(source.Length);
+                        if (source.Length > 0)
+                        {
+                                buffer.View.CopyFromCPU(source);
+                        }
+
+                        return buffer;
+                }
+        }
+
+        private static readonly ConcurrentDictionary<Accelerator, DivisorGeneratorGpuTables> Cache = new(AcceleratorReferenceComparer.Instance);
+
+        public static void WarmUpByDivisorTables(Accelerator accelerator)
+        {
+                if (accelerator is null || accelerator.AcceleratorType == AcceleratorType.CPU)
+                {
+                        return;
+                }
+
+                Cache.GetOrAdd(accelerator, acc => new DivisorGeneratorGpuTables(acc));
+        }
+
+        public static void Clear(Accelerator accelerator)
+        {
+                if (accelerator is null)
+                {
+                        return;
+                }
+
+                if (Cache.TryRemove(accelerator, out var tables))
+                {
+                        tables.Dispose();
+                }
+        }
+
+        public static void DisposeAll()
+        {
+                foreach (var entry in Cache)
+                {
+                        entry.Value.Dispose();
+                }
+
+                Cache.Clear();
+        }
+}