Merge pull request #146 from mlemanczyk/codex/optimize-ulongextensions-for-gpu-calculations

Accelerate Montgomery odd power initialization with GPU batches

PerfectNumbers.Core/Gpu/GpuKernelPool.cs

@@ -37,6 +37,7 @@ public sealed class KernelContainer
             ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>? IncrementalOrder;
     public Action<AcceleratorStream, Index1D, ulong, GpuUInt128, GpuUInt128, byte, ulong,
         ResidueAutomatonArgs, ArrayView<int>, ArrayView1D<ulong, Stride1D.Dense>>? Pow2ModOrder;
+    public Action<AcceleratorStream, Index1D, ArrayView1D<MontgomeryDivisorData, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int>? MontgomeryOddPowers;
 
     // Optional device buffer with small divisor cycles (<= 4M). Index = divisor, value = cycle length.
     public MemoryBuffer1D<ulong, Stride1D.Dense>? SmallCycles;
@@ -189,6 +190,57 @@ public ref struct GpuKernelLease(IDisposable limiter, GpuContextLease gpu, Kerne
         }
     }
 
+    public readonly Action<AcceleratorStream, Index1D, ArrayView1D<MontgomeryDivisorData, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int> MontgomeryOddPowerKernel
+    {
+        get
+        {
+            var accel = Accelerator;
+            lock (_gpu!.Value.KernelInitLock)
+            {
+                return KernelContainer.InitOnce(ref _kernels.MontgomeryOddPowers, () =>
+                {
+                    var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<MontgomeryDivisorData, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int>(MontgomeryOddPowerKernelScan);
+                    var kernel = KernelUtil.GetKernel(loaded);
+                    return kernel.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<MontgomeryDivisorData, Stride1D.Dense>, ArrayView1D<int, Stride1D.Dense>, ArrayView1D<ulong, Stride1D.Dense>, int>>();
+                });
+            }
+        }
+    }
+
+    private static void MontgomeryOddPowerKernelScan(
+        Index1D index,
+        ArrayView1D<MontgomeryDivisorData, Stride1D.Dense> divisors,
+        ArrayView1D<int, Stride1D.Dense> counts,
+        ArrayView1D<ulong, Stride1D.Dense> results,
+        int stride)
+    {
+        int tableIndex = index.X;
+        int baseOffset = tableIndex * stride;
+        MontgomeryDivisorData divisor = divisors[tableIndex];
+        int oddPowerCount = counts[tableIndex];
+        if (oddPowerCount <= 0)
+        {
+            return;
+        }
+
+        ulong montgomeryTwo = divisor.MontgomeryTwo;
+        results[baseOffset] = montgomeryTwo;
+        if (oddPowerCount == 1)
+        {
+            return;
+        }
+
+        ulong modulus = divisor.Modulus;
+        ulong nPrime = divisor.NPrime;
+        ulong square = divisor.MontgomeryTwoSquared;
+        ulong current = montgomeryTwo;
+        for (int i = 1; i < oddPowerCount; i++)
+        {
+            current = ULongExtensions.MontgomeryMultiply(current, square, modulus, nPrime);
+            results[baseOffset + i] = current;
+        }
+    }
+
     // TODO: Plumb the small-cycles device buffer into all kernels that can benefit
     // (some already accept it). Consider a compact type (byte/ushort) for memory footprint.
 

PerfectNumbers.Core/Gpu/MontgomeryOddPowerGpu.cs

@@ -0,0 +1,94 @@
+using System;
+using System.Runtime.InteropServices;
+using ILGPU.Runtime;
+
+namespace PerfectNumbers.Core.Gpu;
+
+internal static class MontgomeryOddPowerGpu
+{
+    public static bool TryCompute(in MontgomeryDivisorData divisor, int oddPowerCount, Span<ulong> destination)
+    {
+        if (oddPowerCount <= 0)
+        {
+            return true;
+        }
+
+        Span<MontgomeryDivisorData> divisorSpan = stackalloc MontgomeryDivisorData[1];
+        divisorSpan[0] = divisor;
+        Span<int> countSpan = stackalloc int[1];
+        countSpan[0] = oddPowerCount;
+        Span<ulong> resultSpan = stackalloc ulong[PerfectNumberConstants.MaxOddPowersCount];
+
+        if (!TryComputeBatch(divisorSpan, countSpan, resultSpan, PerfectNumberConstants.MaxOddPowersCount))
+        {
+            return false;
+        }
+
+        resultSpan[..oddPowerCount].CopyTo(destination);
+        return true;
+    }
+
+    public static bool TryComputeBatch(
+        ReadOnlySpan<MontgomeryDivisorData> divisors,
+        ReadOnlySpan<int> oddPowerCounts,
+        Span<ulong> results,
+        int stride)
+    {
+        int divisorCount = divisors.Length;
+        if (divisorCount == 0)
+        {
+            return true;
+        }
+
+        if (oddPowerCounts.Length < divisorCount)
+        {
+            throw new ArgumentException(
+                "oddPowerCounts span must contain at least as many elements as divisors span.",
+                nameof(oddPowerCounts));
+        }
+
+        if (results.Length < divisorCount * stride)
+        {
+            throw new ArgumentException(
+                "The result span is shorter than the required batched odd power storage.",
+                nameof(results));
+        }
+
+        if (GpuContextPool.ForceCpu)
+        {
+            return false;
+        }
+
+        var lease = GpuKernelPool.GetKernel(useGpuOrder: true);
+        var execution = lease.EnterExecutionScope();
+        var accelerator = lease.Accelerator;
+        var stream = lease.Stream;
+
+        using var divisorBuffer = accelerator.Allocate1D<MontgomeryDivisorData>(divisorCount);
+        using var countBuffer = accelerator.Allocate1D<int>(divisorCount);
+        using var resultBuffer = accelerator.Allocate1D<ulong>(divisorCount * stride);
+
+        try
+        {
+            divisorBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(divisors), divisorCount);
+            countBuffer.View.CopyFromCPU(ref MemoryMarshal.GetReference(oddPowerCounts), divisorCount);
+            resultBuffer.MemSetToZero();
+
+            var kernel = lease.MontgomeryOddPowerKernel;
+            kernel(stream, divisorCount, divisorBuffer.View, countBuffer.View, resultBuffer.View, stride);
+            stream.Synchronize();
+
+            resultBuffer.View.CopyToCPU(ref MemoryMarshal.GetReference(results), divisorCount * stride);
+            return true;
+        }
+        catch
+        {
+            return false;
+        }
+        finally
+        {
+            execution.Dispose();
+            lease.Dispose();
+        }
+    }
+}

PerfectNumbers.Core/ULongExtensions.cs

@@ -535,21 +535,35 @@ public static class ULongExtensions
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private static ulong[] InitializeMontgomeryOddPowersCpu(in MontgomeryDivisorData divisor, ulong modulus, ulong nPrime, int oddPowerCount)
     {
-		ulong[] oddPowers = oddPowerCount < PerfectNumberConstants.PooledArrayThreshold ? new ulong[oddPowerCount] : ArrayPool<ulong>.Shared.Rent(oddPowerCount);
-		
-        oddPowers[0] = divisor.MontgomeryTwo;
+        ulong[] oddPowers = oddPowerCount < PerfectNumberConstants.PooledArrayThreshold
+            ? new ulong[oddPowerCount]
+            : ArrayPool<ulong>.Shared.Rent(oddPowerCount);
+
+        if (oddPowerCount == 0)
+        {
+            return oddPowers;
+        }
+
+        Span<ulong> destination = oddPowers.AsSpan(0, oddPowerCount);
+        bool computedOnGpu = oddPowerCount <= PerfectNumberConstants.MaxOddPowersCount
+            && MontgomeryOddPowerGpu.TryCompute(divisor, oddPowerCount, destination);
+        if (computedOnGpu)
+        {
+            return oddPowers;
+        }
+
+        destination[0] = divisor.MontgomeryTwo;
         if (oddPowerCount == 1)
         {
             return oddPowers;
         }
 
-        ulong previous,
-			  square = divisor.MontgomeryTwoSquared;
-
+        ulong square = divisor.MontgomeryTwoSquared;
+        ulong current = destination[0];
         for (int i = 1; i < oddPowerCount; i++)
         {
-            previous = oddPowers[i - 1];
-            oddPowers[i] = previous.MontgomeryMultiply(square, modulus, nPrime);
+            current = current.MontgomeryMultiply(square, modulus, nPrime);
+            destination[i] = current;
         }
 
         return oddPowers;