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

Reduce branching in Pow2Montgomery GPU execution path
This commit is contained in:
Marek Lemańczyk authored and GitHub committed 2025-10-22 01:43:41 +02:00
commit 4baca1ba70
5 files changed
+109 -52

No files matched your search

@@ -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);
}
}
@@ -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]);
}
}
@@ -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)
@@ -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)
+92 -43
View File
@@ -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;
index -= (int)(currentBit ^ 1UL);
if (currentBit == 0UL)
{
continue;
}
bool processWindow = currentBit != 0UL;
int windowStart = index - windowSize + 1;
// if (windowStart < 0)
result = processWindow ? result : squared;
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;
windowStart &= ~negativeMask;
int negativeMask = windowStartCandidate >> 31;
windowStartCandidate &= ~negativeMask;
windowStart = GetNextSetBitIndexGpu(exponent, windowStart);
int windowLength = index - windowStart + 1;
int windowStart = processWindow ? GetNextSetBitIndexGpu(exponent, windowStartCandidate) : windowStartCandidate;
int windowLength = processWindow ? index - windowStart + 1 : 0;
for (int square = 0; square < windowLength; square++)
{
result = MontgomeryMultiply(result, result, modulus, nPrime);
}
ulong mask = (1UL << windowLength) - 1UL;
ulong windowValue = (exponent >> windowStart) & mask;
ulong multiplier = ComputeMontgomeryOddPowerGpu(windowValue, divisor, modulus, nPrime);
result = MontgomeryMultiply(result, multiplier, modulus, nPrime);
result = processWindow
? MontgomeryMultiply(
result,
ComputeMontgomeryOddPowerGpu(
(exponent >> windowStart) & ((1UL << windowLength) - 1UL),
divisor,
modulus,
nPrime),
modulus,
nPrime)
: result;
index = windowStart - 1;
}
if (!keepMontgomery)
{
result = MontgomeryMultiply(result, 1UL, modulus, nPrime);
index = processWindow ? windowStart - 1 : index;
}
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 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 kernelGroup = KernelCache.GetOrAdd(accelerator, static accel => Pow2MontgomeryKernelGroup.Create(accel));
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.
byte keepFlag = keepMontgomery ? (byte)1 : (byte)0;
var kernel = useKeepKernel ? kernelGroup.KeepMontgomery : kernelGroup.ConvertToStandard;
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);
}
}
}
}