Make GPU mask compaction branchless with shared staging

This commit is contained in:
Marek Lemańczyk committed 2025-10-22 22:01:39 +02:00
1 parent ea0770a84a
commit 4d9158e8f3
2 files changed
+38 -19

No files matched your search

@@ -1,4 +1,5 @@
using ILGPU; using ILGPU;
using ILGPU.Algorithms;
using ILGPU.Runtime; using ILGPU.Runtime;
using PerfectNumbers.Core; using PerfectNumbers.Core;
@@ -179,43 +180,60 @@ internal static class DivisorByDivisorKernels
int groupSize = Group.Dimension.X; int groupSize = Group.Dimension.X;
var shared = SharedMemory.GetDynamic<int>(); var shared = SharedMemory.GetDynamic<int>();
int prefixBase = 0;
int mappingBase = length;
int fallbackBase = length * 2;
// The dynamic shared-memory slice stores prefix sums, compacted index mappings, and fallback slots
// in separate regions so the scatter stage can avoid divergent branches.
byte maskValue = mask[globalIndex]; byte maskValue = mask[globalIndex];
int accepted = maskValue != 0 ? 1 : 0; int accepted = maskValue != 0 ? 1 : 0;
shared[localIndex] = accepted; shared[prefixBase + localIndex] = accepted;
shared[fallbackBase + localIndex] = globalIndex;
ulong candidateValue = sourceCandidates[globalIndex];
Group.Barrier(); Group.Barrier();
int offset = 1; int offset = 1;
while (offset < groupSize) while (offset < groupSize)
{ {
int addend = localIndex >= offset ? shared[localIndex - offset] : 0; int addendIndex = prefixBase + localIndex - offset;
int addend = localIndex >= offset ? shared[addendIndex] : 0;
Group.Barrier(); Group.Barrier();
int current = shared[localIndex]; int currentIndex = prefixBase + localIndex;
int current = shared[currentIndex];
int sum = current + addend; int sum = current + addend;
if (localIndex >= offset) int shouldUpdate = localIndex >= offset ? 1 : 0;
{ shared[currentIndex] = shouldUpdate != 0 ? sum : current;
shared[localIndex] = sum;
}
Group.Barrier(); Group.Barrier();
offset <<= 1; offset <<= 1;
} }
if (accepted != 0) int inclusivePrefix = shared[prefixBase + localIndex];
{ int exclusiveIndex = inclusivePrefix - accepted;
int targetIndex = shared[localIndex] - 1; int targetBase = accepted != 0 ? mappingBase : fallbackBase;
compactedCandidates[targetIndex] = candidateValue; int targetOffset = accepted != 0 ? exclusiveIndex : localIndex;
} shared[targetBase + targetOffset] = globalIndex;
if (globalIndex == length - 1) Group.Barrier();
{
compactedCount[0] = shared[localIndex]; int totalAccepted = shared[prefixBase + length - 1];
} int writerMask = localIndex < totalAccepted ? 1 : 0;
int lookupBase = writerMask != 0 ? mappingBase : fallbackBase;
int sourceIndex = shared[lookupBase + localIndex];
ulong candidateValue = sourceCandidates[sourceIndex];
int destinationIndex = localIndex;
ulong existingValue = compactedCandidates[destinationIndex];
compactedCandidates[destinationIndex] = writerMask != 0 ? candidateValue : existingValue;
int lastLaneMask = globalIndex == length - 1 ? 1 : 0;
int countValue = shared[prefixBase + localIndex];
int storedCount = lastLaneMask != 0 ? countValue : 0;
Atomic.Max(ref compactedCount[0], storedCount);
} }
public static void ComputeMontgomeryExponentKernel(Index1D index, MontgomeryDivisorData divisor, ArrayView<ulong> exponents, ArrayView<ulong> results) public static void ComputeMontgomeryExponentKernel(Index1D index, MontgomeryDivisorData divisor, ArrayView<ulong> exponents, ArrayView<ulong> results)
@@ -588,7 +588,8 @@ public sealed class MersenneNumberDivisorByDivisorGpuTester : IMersenneNumberDiv
// return 0; // return 0;
// } // }
SharedMemoryConfig sharedMemoryConfig = SharedMemoryConfig.RequestDynamic<int>(count); countView.MemSetToZero();
SharedMemoryConfig sharedMemoryConfig = SharedMemoryConfig.RequestDynamic<int>(count * 3);
Index1D gridDim = new Index1D(1); Index1D gridDim = new Index1D(1);
Index1D groupDim = new Index1D(count); Index1D groupDim = new Index1D(count);
KernelConfig kernelConfig = new KernelConfig(gridDim, groupDim, sharedMemoryConfig); KernelConfig kernelConfig = new KernelConfig(gridDim, groupDim, sharedMemoryConfig);