Merge pull request #223 from mlemanczyk/codex/refactor-pow2montgomerygpuexecutor.execute-for-branch-less-c

Reduce branching in Pow2Montgomery GPU execution path

PerfectNumbers.Core/Gpu/Kernels/DivisorByDivisorKernels.cs

@@ -17,7 +17,7 @@ internal static class DivisorByDivisorKernels
         }
 
         ulong exponent = exponents[index];
-        hits[index] = ULongExtensions.Pow2MontgomeryModWindowedGpu(divisor, exponent, keepMontgomery: false) == 1UL ? (byte)1 : (byte)0;
+        hits[index] = ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisor, exponent) == 1UL ? (byte)1 : (byte)0;
     }
 
     public static void ComputeRemainderDeltasKernel(Index1D index, ArrayView<ulong> gaps, ArrayView<byte> deltas, byte modulus)
@@ -149,6 +149,6 @@ internal static class DivisorByDivisorKernels
 
         ulong exponent = exponents[index];
 
-        results[index] = ULongExtensions.Pow2MontgomeryModWindowedGpu(divisor, exponent, keepMontgomery: true);
+        results[index] = ULongExtensions.Pow2MontgomeryModWindowedGpuKeepMontgomery(divisor, exponent);
     }
 }

PerfectNumbers.Core/Gpu/Kernels/Pow2MontgomeryKernels.cs

@@ -6,13 +6,21 @@ namespace PerfectNumbers.Core.Gpu;
 
 internal static class Pow2MontgomeryKernels
 {
-    public static void Pow2MontgomeryKernel(
+    public static void Pow2MontgomeryKernelKeepMontgomery(
         Index1D index,
         ArrayView1D<ulong, Stride1D.Dense> exponents,
         MontgomeryDivisorData divisor,
-        byte keepMontgomery,
         ArrayView1D<ulong, Stride1D.Dense> results)
     {
-        results[index] = ULongExtensions.Pow2MontgomeryModWindowedGpu(divisor, exponents[index], keepMontgomery != 0);
+        results[index] = ULongExtensions.Pow2MontgomeryModWindowedGpuKeepMontgomery(divisor, exponents[index]);
+    }
+
+    public static void Pow2MontgomeryKernelConvertToStandard(
+        Index1D index,
+        ArrayView1D<ulong, Stride1D.Dense> exponents,
+        MontgomeryDivisorData divisor,
+        ArrayView1D<ulong, Stride1D.Dense> results)
+    {
+        results[index] = ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisor, exponents[index]);
     }
 }

PerfectNumbers.Core/Gpu/Kernels/PrimeOrderKernels.cs

@@ -347,7 +347,7 @@ internal static partial class PrimeOrderGpuHeuristics
             }
 
             ulong reduced = phi / factor;
-            if (ULongExtensions.Pow2MontgomeryModWindowedGpu(divisor, reduced, keepMontgomery: false) == 1UL)
+            if (ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisor, reduced) == 1UL)
             {
                 return false;
             }
@@ -1131,7 +1131,7 @@ internal static partial class PrimeOrderGpuHeuristics
 
     private static bool Pow2EqualsOneKernel(ulong exponent, in MontgomeryDivisorData divisor)
     {
-        return ULongExtensions.Pow2MontgomeryModWindowedGpu(divisor, exponent, keepMontgomery: false) == 1UL;
+        return ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisor, exponent) == 1UL;
     }
 
     private static ulong CalculateByDoublingKernel(ulong prime)
@@ -1156,7 +1156,7 @@ internal static partial class PrimeOrderGpuHeuristics
     internal static void Pow2ModKernel(Index1D index, ArrayView1D<ulong, Stride1D.Dense> exponents, MontgomeryDivisorData divisor, ArrayView1D<ulong, Stride1D.Dense> remainders)
     {
         ulong exponent = exponents[index];
-        remainders[index] = ULongExtensions.Pow2MontgomeryModWindowedGpu(divisor, exponent, keepMontgomery: false);
+        remainders[index] = ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisor, exponent);
     }
 
     private static Action<AcceleratorStream, Index1D, ArrayView1D<GpuUInt128, Stride1D.Dense>, GpuUInt128, ArrayView1D<GpuUInt128, Stride1D.Dense>> GetPow2ModWideKernel(Accelerator accelerator)

PerfectNumbers.Core/Gpu/PrimeOrderGpuHeuristics.cs

@@ -652,7 +652,7 @@ internal static partial class PrimeOrderGpuHeuristics
             return 0UL;
         }
 
-        return ULongExtensions.Pow2MontgomeryModWindowedGpu(divisorData, exponent, keepMontgomery: false);
+        return ULongExtensions.Pow2MontgomeryModWindowedGpuConvertToStandard(divisorData, exponent);
     }
 
     private static Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> GetPow2ModKernel(Accelerator accelerator)

PerfectNumbers.Core/ULongExtensions.Gpu.cs

@@ -109,20 +109,35 @@ public static partial class ULongExtensions
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        internal static ulong Pow2MontgomeryModWindowedGpu(in MontgomeryDivisorData divisor, ulong exponent, bool keepMontgomery)
+        internal static ulong Pow2MontgomeryModWindowedGpuKeepMontgomery(in MontgomeryDivisorData divisor, ulong exponent)
+        {
+                return Pow2MontgomeryModWindowedGpuMontgomeryResult(divisor, exponent);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static ulong Pow2MontgomeryModWindowedGpuConvertToStandard(in MontgomeryDivisorData divisor, ulong exponent)
+        {
+                ulong modulus = divisor.Modulus;
+                ulong nPrime = divisor.NPrime;
+                ulong montgomeryResult = Pow2MontgomeryModWindowedGpuMontgomeryResult(divisor, exponent);
+                return MontgomeryMultiply(montgomeryResult, 1UL, modulus, nPrime);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static ulong Pow2MontgomeryModWindowedGpuMontgomeryResult(in MontgomeryDivisorData divisor, ulong exponent)
         {
                 ulong modulus = divisor.Modulus;
                 // Reuse the optimized MontgomeryMultiply extension highlighted in MontgomeryMultiplyBenchmarks so the GPU path matches the fastest CPU reduction.
                 // We shouldn't hit it in production code. If we do, we're doing something wrong.
                 // if (exponent == 0UL)
                 // {
-                //         return keepMontgomery ? divisor.MontgomeryOne : 1UL % modulus;
+                //         return divisor.MontgomeryOne;
                 // }
 
                 // We barely ever hit it in production code. It's e.g. 772 calls out of billions
                 // if (exponent <= Pow2WindowFallbackThreshold)
                 // {
-                //         return Pow2MontgomeryModSingleBit(exponent, divisor, keepMontgomery);
+                //         return Pow2MontgomeryModSingleBit(exponent, divisor);
                 // }
 
                 int bitLength = GetPortableBitLengthGpu(exponent);
@@ -135,42 +150,41 @@ public static partial class ULongExtensions
                 {
                         ulong currentBit = (exponent >> index) & 1UL;
                         ulong squared = MontgomeryMultiply(result, result, modulus, nPrime);
-                        result = currentBit == 0UL ? squared : result;
+                        bool processWindow = currentBit != 0UL;
-                        index -= (int)(currentBit ^ 1UL);
-                        if (currentBit == 0UL)
-                        {
-                                continue;
-                        }
 
-                        int windowStart = index - windowSize + 1;
+                        result = processWindow ? result : squared;
-                        // if (windowStart < 0)
+                        index -= (int)(currentBit ^ 1UL);
+
+                        int windowStartCandidate = index - windowSize + 1;
+                        // if (windowStartCandidate < 0)
                         // {
-                        //         windowStart = 0;
+                        //         windowStartCandidate = 0;
                         // }
                         // Clamp the negative offset without branching so the GPU loop stays divergence-free.
                         // This still handles windows that would otherwise extend past the most significant bit.
-                        int negativeMask = windowStart >> 31;
+                        int negativeMask = windowStartCandidate >> 31;
-                        windowStart &= ~negativeMask;
+                        windowStartCandidate &= ~negativeMask;
 
-                        windowStart = GetNextSetBitIndexGpu(exponent, windowStart);
+                        int windowStart = processWindow ? GetNextSetBitIndexGpu(exponent, windowStartCandidate) : windowStartCandidate;
-
+                        int windowLength = processWindow ? index - windowStart + 1 : 0;
-                        int windowLength = index - windowStart + 1;
                         for (int square = 0; square < windowLength; square++)
                         {
                                 result = MontgomeryMultiply(result, result, modulus, nPrime);
                         }
 
-                        ulong mask = (1UL << windowLength) - 1UL;
+                        result = processWindow
-                        ulong windowValue = (exponent >> windowStart) & mask;
+                                ? MontgomeryMultiply(
-                        ulong multiplier = ComputeMontgomeryOddPowerGpu(windowValue, divisor, modulus, nPrime);
+                                        result,
-                        result = MontgomeryMultiply(result, multiplier, modulus, nPrime);
+                                        ComputeMontgomeryOddPowerGpu(
+                                                (exponent >> windowStart) & ((1UL << windowLength) - 1UL),
+                                                divisor,
+                                                modulus,
+                                                nPrime),
+                                        modulus,
+                                        nPrime)
+                                : result;
 
-                        index = windowStart - 1;
+                        index = processWindow ? windowStart - 1 : index;
-                }
-
-                if (!keepMontgomery)
-                {
-                        result = MontgomeryMultiply(result, 1UL, modulus, nPrime);
                 }
 
                 return result;
@@ -183,16 +197,18 @@ public static partial class ULongExtensions
                 // ulong modulus = divisor.Modulus;
                 // if (exponent == 0UL)
                 // {
-                //      return keepMontgomery ? divisor.MontgomeryOne : 1UL % modulus;
+                //      return divisor.MontgomeryOne;
                 // }
 
                 // We barely ever hit it in production code. It's e.g. 772 calls out of billions
                 // if (exponent <= Pow2WindowFallbackThreshold)
                 // {
-                //      return Pow2MontgomeryModSingleBit(exponent, divisor, keepMontgomery);
+                //      return Pow2MontgomeryModSingleBit(exponent, divisor);
                 // }
 
-                return Pow2MontgomeryGpuExecutor.Execute(exponent, divisor, keepMontgomery);
+                return keepMontgomery
+                        ? Pow2MontgomeryGpuExecutor.ExecuteKeep(exponent, divisor)
+                        : Pow2MontgomeryGpuExecutor.ExecuteConvert(exponent, divisor);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -289,20 +305,30 @@ public static partial class ULongExtensions
         public static ulong Pow2MontgomeryModWithCycleGpu(this ulong exponent, ulong cycleLength, in MontgomeryDivisorData divisor)
         {
                 ulong rotationCount = exponent % cycleLength;
-                return Pow2MontgomeryModWindowedGpu(rotationCount, divisor, keepMontgomery: false);
+                return Pow2MontgomeryGpuExecutor.ExecuteConvert(rotationCount, divisor);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static ulong Pow2MontgomeryModFromCycleRemainderGpu(this ulong reducedExponent, in MontgomeryDivisorData divisor)
         {
-                return Pow2MontgomeryModWindowedGpu(reducedExponent, divisor, keepMontgomery: false);
+                return Pow2MontgomeryGpuExecutor.ExecuteConvert(reducedExponent, divisor);
         }
 
         private static class Pow2MontgomeryGpuExecutor
         {
-                private static readonly ConcurrentDictionary<Accelerator, Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, byte, ArrayView1D<ulong, Stride1D.Dense>>> KernelCache = new();
+                private static readonly ConcurrentDictionary<Accelerator, Pow2MontgomeryKernelGroup> KernelCache = new();
 
-                public static ulong Execute(ulong exponent, in MontgomeryDivisorData divisor, bool keepMontgomery)
+                public static ulong ExecuteKeep(ulong exponent, in MontgomeryDivisorData divisor)
+                {
+                        return Execute(exponent, divisor, useKeepKernel: true);
+                }
+
+                public static ulong ExecuteConvert(ulong exponent, in MontgomeryDivisorData divisor)
+                {
+                        return Execute(exponent, divisor, useKeepKernel: false);
+                }
+
+                private static ulong Execute(ulong exponent, in MontgomeryDivisorData divisor, bool useKeepKernel)
                 {
                         ulong result = 0UL;
 
@@ -316,12 +342,7 @@ public static partial class ULongExtensions
                         //         return 0UL;
                         // }
 
-                        var kernel = KernelCache.GetOrAdd(accelerator, static accel =>
+                        var kernelGroup = KernelCache.GetOrAdd(accelerator, static accel => Pow2MontgomeryKernelGroup.Create(accel));
-                        {
-                                var loaded = accel.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, byte, ArrayView1D<ulong, Stride1D.Dense>>(Pow2MontgomeryKernels.Pow2MontgomeryKernel);
-                                var launcher = KernelUtil.GetKernel(loaded);
-                                return launcher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, byte, ArrayView1D<ulong, Stride1D.Dense>>>();
-                        });
 
                         var exponentBuffer = accelerator.Allocate1D<ulong>(1);
                         var resultBuffer = accelerator.Allocate1D<ulong>(1);
@@ -330,10 +351,9 @@ public static partial class ULongExtensions
                         // We don't need to worry about any left-overs here.
                         // resultBuffer.MemSetToZero();
 
-                        // TODO: Have 2 kernels for both cases to have branch-less solution.
+                        var kernel = useKeepKernel ? kernelGroup.KeepMontgomery : kernelGroup.ConvertToStandard;
-                        byte keepFlag = keepMontgomery ? (byte)1 : (byte)0;
                         AcceleratorStream stream = lease.Stream;
-                        kernel(stream, 1, exponentBuffer.View, divisor, keepFlag, resultBuffer.View);
+                        kernel(stream, 1, exponentBuffer.View, divisor, resultBuffer.View);
                         stream.Synchronize();
 
                         resultBuffer.View.CopyToCPU(ref result, 1);
@@ -354,5 +374,34 @@ public static partial class ULongExtensions
                         // Intentionally avoid exception handling here; any accelerator failure should crash the scanner.
                         return result;
                 }
+
+                private sealed class Pow2MontgomeryKernelGroup
+                {
+                        internal Pow2MontgomeryKernelGroup(
+                                Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> keepMontgomery,
+                                Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> convertToStandard)
+                        {
+                                KeepMontgomery = keepMontgomery;
+                                ConvertToStandard = convertToStandard;
+                        }
+
+                        internal Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> KeepMontgomery { get; }
+
+                        internal Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>> ConvertToStandard { get; }
+
+                        internal static Pow2MontgomeryKernelGroup Create(Accelerator accelerator)
+                        {
+                                var keepKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2MontgomeryKernels.Pow2MontgomeryKernelKeepMontgomery);
+                                var keepLauncher = KernelUtil.GetKernel(keepKernel);
+                                var keepDelegate = keepLauncher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
+
+                                var convertKernel = accelerator.LoadAutoGroupedStreamKernel<Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>(Pow2MontgomeryKernels.Pow2MontgomeryKernelConvertToStandard);
+                                var convertLauncher = KernelUtil.GetKernel(convertKernel);
+                                var convertDelegate = convertLauncher.CreateLauncherDelegate<Action<AcceleratorStream, Index1D, ArrayView1D<ulong, Stride1D.Dense>, MontgomeryDivisorData, ArrayView1D<ulong, Stride1D.Dense>>>();
+
+                                return new Pow2MontgomeryKernelGroup(keepDelegate, convertDelegate);
+                        }
+                }
+
         }
 }