@woosh/meep-engine 3.2.0 → 3.4.0

package/src/shade/renderer/gpu_primitive/sort/one_sweep/shader_digit_binning_pass_pairs.js

@@ -1,990 +0,0 @@
-import { WebGPUExtensionType } from "../../../../descriptor/WebGPUExtensionType.js";
-import { atomic_u32, u32 } from "../../../shader/type/WEBGPU_WGSL_PRIMITIVE_TYPES.js";
-import { WebGPUArray } from "../../../shader/type/WebGPUArray.js";
-import { WebGPUStruct } from "../../../shader/type/WebGPUStruct.js";
-import { CodeChunk } from "../../../shader/compiler/CodeChunk.js";
-import { ComputeShader } from "../../../shader/ComputeShader.js";
-import { ShaderResourceSetDescriptor } from "../../../shader/resource/ShaderResourceSetDescriptor.js";
-import { k_keysPerThread, k_partitionSize, k_radix } from "./defs.js";
-
-export const PUSH_CONSTANTS = WebGPUStruct.from({
-    e_radixShift: u32,
-    e_numKeys: u32,
-    e_threadBlocks: u32
-});
-
-const resources = new ShaderResourceSetDescriptor();
-resources.createGroup()
-    .addStorageBuffer("b_index", WebGPUArray.from(atomic_u32), true)
-    .addStorageBuffer("b_sort", WebGPUArray.u32)
-    .addStorageBuffer("b_sortPayload", WebGPUArray.u32)
-    .addStorageBuffer("b_alt", WebGPUArray.u32, true)
-    .addStorageBuffer("b_altPayload", WebGPUArray.u32, true)
-    .addStorageBuffer("b_passHist", WebGPUArray.from(atomic_u32), true)
-    .addUniform("settings", PUSH_CONSTANTS);
-
-/**
- * Based off of Research by:
- *          Andy Adinets, Nvidia Corporation
- *          Duane Merrill, Nvidia Corporation
- *          https://research.nvidia.com/publication/2022-06_onesweep-faster-least-significant-digit-radix-sort-gpus
- *
- * @type {CodeChunk}
- */
-const body = CodeChunk.from(
-    //language=WGSL
-    `
-const N_DIM             = 256;
-const D_DIM             = 256u;
-const PART_SIZE         = ${k_partitionSize};
-const D_TOTAL_SMEM      = 4096u;
-const KEYS_PER_THREAD   = ${k_keysPerThread};
-
-const RADIX             = ${k_radix};    //Number of digit bins
-const RADIX_MASK        = ${k_radix-1};
-const HALF_RADIX        = 128u;    //For smaller waves where bit packing is necessary
-const HALF_MASK         = 127u;
-const RADIX_LOG:u32     = ${Math.log2(k_radix)};
-
-const FLAG_REDUCTION:u32      = 1;       //Flag value indicating reduction of a partition tile is ready
-const FLAG_INCLUSIVE:u32      = 2;       //Flag value indicating inclusive sum of a partition tile is ready
-const FLAG_MASK:u32           = 3;       //Mask used to retrieve flag values
-
-
-// private (not workgroup): g_d below already consumes the ENTIRE 16384-byte default workgroup
-// budget, and every invocation writes the identical builtin value anyway.
-var<private> subgroup_size:u32;
-var<private> subgroup_index: u32;
-var<workgroup> g_d: array<atomic<u32>,D_TOTAL_SMEM>; //Shared memory for DigitBinningPass and DownSweep kernels
-
-struct KeyStruct
-{
-    k:array<u32,KEYS_PER_THREAD>
-};
-
-struct OffsetStruct
-{
-    o:array<u32,KEYS_PER_THREAD>
-};
-
-struct DigitStruct
-{
-    d:array<u32,KEYS_PER_THREAD>
-};
-
-fn ExtractDigit(key:u32) -> u32
-{
-    return (key >> settings.e_radixShift) & RADIX_MASK;
-}
-
-fn CurrentPass() -> u32
-{
-    return settings.e_radixShift >> 3;
-}
-fn PassHistOffset(index:u32) -> u32
-{
-    return ((CurrentPass() * settings.e_threadBlocks) + index) << RADIX_LOG;
-}
-
-fn ExtractPackedIndex(key:u32) -> u32
-{
-    return (key >> (settings.e_radixShift + 1)) & HALF_MASK;
-}
-
-fn ExtractPackedShift(key:u32) -> u32
-{
-    return select(0u, 16u, ( (key >> settings.e_radixShift) & 1 ) != 0);
-}
-
-fn ExtractPackedValue(packed:u32, key:u32) -> u32
-{
-    return ( packed >> ExtractPackedShift(key) ) & 0xffff;
-}
-
-
-
-fn WaveHistsSizeWGE16() -> u32
-{
-    return D_DIM / subgroup_size * RADIX;
-}
-
-fn WaveHistsSizeWLT16() -> u32
-{
-    return D_TOTAL_SMEM;
-}
-
-fn SerialIterations() -> u32
-{
-    return (D_DIM / subgroup_size + 31) >> 5;
-}
-
-fn ClearWaveHists( gtid:u32 )
-{
-    let histsEnd = select(WaveHistsSizeWLT16(),WaveHistsSizeWGE16(),subgroup_size >= 16);
-        
-    for (var i = gtid; i < histsEnd; i += D_DIM){
-        atomicStore(&g_d[i], 0);
-    }
-}
-
-//*****************************************************************************
-//HELPER FUNCTIONS
-//*****************************************************************************
-fn getWaveIndex(gtid:u32) -> u32
-{
-    return gtid / subgroup_size;
-}
-
-fn SubPartSizeWGE16() -> u32
-{
-    return KEYS_PER_THREAD * subgroup_size;
-}
-
-fn SharedOffsetWGE16( gtid:u32) -> u32
-{
-    return subgroup_index + getWaveIndex(gtid) * SubPartSizeWGE16();
-}
-
-fn SubPartSizeWLT16( _serialIterations:u32) -> u32
-{
-    return KEYS_PER_THREAD * subgroup_size * _serialIterations;
-}
-
-fn SharedOffsetWLT16( gtid:u32 , _serialIterations: u32) -> u32
-{
-    return subgroup_index +
-        (getWaveIndex(gtid) / _serialIterations * SubPartSizeWLT16(_serialIterations)) +
-        (getWaveIndex(gtid) % _serialIterations * subgroup_size);
-}
-
-fn DeviceOffsetWGE16(gtid:u32, partIndex:u32) -> u32
-{
-    return SharedOffsetWGE16(gtid) + partIndex * PART_SIZE;
-}
-
-fn DeviceOffsetWLT16(gtid:u32, partIndex:u32, serialIterations:u32) -> u32
-{
-    return SharedOffsetWLT16(gtid, serialIterations) + partIndex * PART_SIZE;
-}
-
-//*****************************************************************************
-//FUNCTIONS COMMON TO THE DOWNSWEEP / DIGIT BINNING PASS
-//*****************************************************************************
-//If the size of  a wave is too small, we do not have enough space in
-//shared memory to assign a histogram to each wave, so instead,
-//some operations are peformed serially.
-fn LoadDummyKey( key: ptr<function,u32> )
-{
-    *key = 0xffffffff;
-}
-
-fn LoadKey( key: ptr<function,u32>,  index: u32)
-{
-    *key = b_sort[index];
-}
-
-fn LoadKeysWGE16(gtid:u32, partIndex:u32) -> KeyStruct
-{
-    var keys:KeyStruct;
-    
-    var t = DeviceOffsetWGE16(gtid, partIndex);
-    for (var i:u32 = 0; 
-        i < KEYS_PER_THREAD;
-        i++
-        ){
-        
-        LoadKey(&keys.k[i], t);
-        t += subgroup_size;
-        
-    }
-    
-    return keys;
-}
-
-fn LoadKeysWLT16( 
-    gtid:u32,
-    partIndex:u32,
-    serialIterations:u32
-) -> KeyStruct {
-    var keys:KeyStruct;
-    
-    var t = DeviceOffsetWLT16(gtid, partIndex, serialIterations);
-    for (var i:u32 = 0; 
-        i < KEYS_PER_THREAD;
-        i++
-        ){
-        
-        LoadKey(&keys.k[i], t);
-        t += subgroup_size * serialIterations;
-        
-    }
-    
-    return keys;
-}
-
-fn LoadKeysPartialWGE16(gtid:u32, partIndex:u32) -> KeyStruct
-{
-    var keys:KeyStruct;
-    
-    var t = DeviceOffsetWGE16(gtid, partIndex);
-    for (var i:u32 = 0; 
-        i < KEYS_PER_THREAD;
-        i++
-        ){
-        
-        if(t < settings.e_numKeys){
-            LoadKey(&keys.k[i], t);
-        }else{
-            LoadDummyKey(&keys.k[i]);
-        }
-        
-        t += subgroup_size;
-        
-    }
-    
-    return keys;
-}
-
-fn LoadKeysPartialWLT16( gtid:u32, partIndex:u32, serialIterations:u32) -> KeyStruct
-{
-    var keys:KeyStruct;
-    
-    var t = DeviceOffsetWLT16(gtid, partIndex, serialIterations);
-    for (var i:u32 = 0; 
-        i < KEYS_PER_THREAD;
-        i++
-        ){
-        
-         if(t < settings.e_numKeys){
-            LoadKey(&keys.k[i], t);
-        }else{
-            LoadDummyKey(&keys.k[i]);
-        }
-        
-        t += subgroup_size * serialIterations;
-        
-    }
-    
-    return keys;
-}
-fn WaveFlagsWGE16() -> u32 {
-    if (subgroup_size == 32u) {
-        // note, using explicit condition to keep WebGPU compiler happy
-        return 0xffffffffu;
-    } else {
-        return (1u << subgroup_size) - 1;
-    }
-}
-
-fn WaveFlagsWLT16() -> u32 {
-    if (subgroup_size == 32u) {
-        return 0xffffffffu; // This case is unreachable, but prevents the compilation error.
-    } else {
-        return (1u << subgroup_size) - 1;
-    }
-}
-
-fn WarpLevelMultiSplitWGE16( key:u32, waveParts:u32, waveFlags:ptr<function,vec4u> )
-{
-    for (var k = 0u; k < RADIX_LOG; k++)
-    {
-        let t:bool = bool( (key >> (k + settings.e_radixShift)) & 1 );
-        let ballot = subgroupBallot(t);
-        
-        for (var wavePart = 0u; wavePart < waveParts; wavePart++){
-            waveFlags[wavePart] &= select( 0xffffffffu, 0u, t) ^ ballot[wavePart];
-        }
-    }
-}
-
-fn WarpLevelMultiSplitWLT16(key:u32, waveFlags:ptr<function,u32>)
-{
-    for (var k = 0u; k < RADIX_LOG; k++)
-    {
-        let t:bool = bool( ( key >> (k + settings.e_radixShift) ) & 1 );
-        *waveFlags &= select(0xffffffffu, 0u, t) ^ subgroupBallot(t).x;
-    }
-}
-
-fn CountPeerBits(
-    peerBits: ptr<function, u32>,
-    totalBits: ptr<function, u32>,
-    waveFlags: vec4u,
-    waveParts: u32
-){
-    for (var wavePart = 0u; wavePart < waveParts; wavePart++)
-    {
-        if (subgroup_index >= wavePart * 32)
-        {
-            let ltMask = select((1u << (subgroup_index & 31)) - 1,0xffffffff,subgroup_index >= (wavePart + 1) * 32 );
-            *peerBits += countOneBits(waveFlags[wavePart] & ltMask);
-        }
-        *totalBits += countOneBits(waveFlags[wavePart]);
-    }
-}
-
-fn CountPeerBitsWLT16(
-    waveFlags: u32,
-    ltMask: u32
-) -> u32 {
-    return countOneBits(waveFlags & ltMask);
-}
-
-fn FindLowestRankPeer(
-    waveFlags: vec4u,
-    waveParts: u32
-) -> u32 {
-    var lowestRankPeer:u32 = 0;
-    for (var wavePart:u32 = 0; wavePart < waveParts; wavePart++)
-    {
-        let fbl = countTrailingZeros( waveFlags[wavePart] );
-        if (fbl == 0xffffffff){
-            lowestRankPeer += 32;
-        } else{
-            return lowestRankPeer + fbl;
-        }
-    }
-    return 0; //will never happen
-}
-
-fn RankKeysWGE16( gtid:u32, keys:KeyStruct) -> OffsetStruct
-{
-    var offsets:OffsetStruct;
-    let waveParts = (subgroup_size + 31) / 32;
-
-    for (var i:u32 = 0; i < KEYS_PER_THREAD; i++)
-    {
-        var waveFlags = vec4(WaveFlagsWGE16());
-        WarpLevelMultiSplitWGE16(keys.k[i], waveParts, &waveFlags);
-        
-        let index = ExtractDigit(keys.k[i]) + (getWaveIndex(gtid) * RADIX);
-        let lowestRankPeer = FindLowestRankPeer(waveFlags, waveParts);
-        
-        var peerBits = 0u;
-        var totalBits = 0u;
-        CountPeerBits(&peerBits, &totalBits, waveFlags, waveParts);
-        
-        var preIncrementVal:u32;
-        
-        if (peerBits == 0){
-            preIncrementVal = atomicAdd(&g_d[index], totalBits);
-        }
-        
-        offsets.o[i] = subgroupShuffle(preIncrementVal, lowestRankPeer) + peerBits;
-    }
-    
-    return offsets;
-}
-
-fn RankKeysWLT16( gtid:u32, keys: KeyStruct , serialIterations:u32) -> OffsetStruct
-{
-    var offsets:OffsetStruct;
-    let ltMask:u32 = (1u << subgroup_index) - 1;
-    
-    for (var i:u32 = 0; i < KEYS_PER_THREAD; i++)
-    {
-        var waveFlags = WaveFlagsWLT16();
-        WarpLevelMultiSplitWLT16(keys.k[i], &waveFlags);
-        
-        let index = ExtractPackedIndex(keys.k[i]) +
-                (getWaveIndex(gtid) / serialIterations * HALF_RADIX);
-        
-        let peerBits = CountPeerBitsWLT16(waveFlags, ltMask);
-        for (var k:u32 = 0; k < serialIterations; k++)
-        {
-            if (getWaveIndex(gtid) % serialIterations == k){
-                offsets.o[i] = ExtractPackedValue(atomicLoad(&g_d[index]), keys.k[i]) + peerBits;
-            }
-            
-            workgroupBarrier();
-            
-            if (getWaveIndex(gtid) % serialIterations == k && peerBits == 0)
-            {
-                atomicAdd(&g_d[index],
-                    countOneBits(waveFlags) << ExtractPackedShift(keys.k[i]));
-            }
-            
-            workgroupBarrier();
-        }
-    }
-    
-    return offsets;
-}
-
-fn WaveHistInclusiveScanCircularShiftWGE16( gtid: u32 ) -> u32
-{
-    var histReduction = atomicLoad(&g_d[gtid]);
-    for (var i:u32 = gtid + RADIX; i < WaveHistsSizeWGE16(); i += RADIX)
-    {
-        let a = atomicLoad(&g_d[i]);
-        histReduction += a;
-        atomicStore(&g_d[i], histReduction - a);
-    }
-    return histReduction;
-}
-
-fn WaveHistInclusiveScanCircularShiftWLT16( gtid: u32 ) -> u32 {
-    var histReduction = atomicLoad(&g_d[gtid]);
-    for (var i:u32 = gtid + HALF_RADIX; i < WaveHistsSizeWLT16(); i += HALF_RADIX)
-    {
-        let a = atomicLoad(&g_d[i]);
-        histReduction += a;
-        atomicStore(&g_d[i], histReduction - a);
-    }
-    return histReduction;
-}
-
-fn WaveHistReductionExclusiveScanWGE16( gtid:u32, histReduction:u32)
-{
-    if (gtid < RADIX){
-        let laneMask = subgroup_size - 1;
-        let address = ((subgroup_index + 1) & laneMask) + (gtid & ~laneMask);
-        atomicStore(&g_d[address], histReduction);
-    }
-    workgroupBarrier();
-    
-    let a0 = subgroupExclusiveAdd(atomicLoad(&g_d[gtid * subgroup_size]));            
-    if (gtid < RADIX / subgroup_size){
-        atomicStore(&g_d[gtid * subgroup_size], a0);
-    }
-    workgroupBarrier();
-
-    let a1 = subgroupBroadcast(atomicLoad(&g_d[gtid - 1]), 1);
-    if (gtid < RADIX && subgroup_index != 0){
-        atomicAdd(&g_d[gtid], a1);
-    }
-}
-
-//inclusive/exclusive prefix sum up the histograms,
-//use a blelloch scan for in place packed exclusive
-fn WaveHistReductionExclusiveScanWLT16( gtid:u32 )
-{
-    var shift:u32 = 1;
-    for (var j:u32 = RADIX >> 2; j > 0; j >>= 1)
-    {
-        workgroupBarrier();
-        if (gtid < j)
-        {
-            let a = ((((gtid << 1) + 2) << shift) - 1) >> 1;
-            let b = ((((gtid << 1) + 1) << shift) - 1) >> 1;
-            atomicAdd(&g_d[a],
-                atomicLoad(&g_d[b]) & 0xffff0000);
-        }
-        shift++;
-    }
-    workgroupBarrier();
-                
-    if (gtid == 0){
-        atomicAnd(&g_d[HALF_RADIX - 1], 0xffff);
-    }
-                
-    for (var j:u32 = 1; j < HALF_RADIX; j <<= 1)
-    {
-        shift--;
-        workgroupBarrier();
-        if (gtid < j)
-        {
-            let t = ((((gtid << 1) + 1) << shift) - 1) >> 1;
-            let t2 = ((((gtid << 1) + 2) << shift) - 1) >> 1;
-            let t3 = atomicLoad(&g_d[t]);
-            atomicStore(&g_d[t], (t3 & 0xffff) | (atomicLoad(&g_d[t2]) & 0xffff0000));
-            atomicAdd(&g_d[t2], t3 & 0xffff0000);
-        }
-    }
-
-    workgroupBarrier();
-    if (gtid < HALF_RADIX)
-    {
-        let t = atomicLoad(&g_d[gtid]);
-        atomicStore(&g_d[gtid], (t >> 16) + (t << 16) + (t & 0xffff0000));
-    }
-}
-
-fn UpdateOffsetsWGE16(gtid:u32, offsets:ptr<function,OffsetStruct>, keys:KeyStruct)
-{
-    if (gtid >= subgroup_size)
-    {
-        let t = getWaveIndex(gtid) * RADIX;
-        for (var i:u32 = 0; i < KEYS_PER_THREAD; i++)
-        {
-            let t2 = ExtractDigit(keys.k[i]);
-            offsets.o[i] += atomicLoad(&g_d[t2 + t]) + atomicLoad(&g_d[t2]);
-        }
-    }
-    else{
-        for (var i:u32 = 0; i < KEYS_PER_THREAD; i++){
-            offsets.o[i] += atomicLoad(&g_d[ExtractDigit(keys.k[i])]);
-        }
-    }
-}
-
-fn UpdateOffsetsWLT16(
-    gtid:u32,
-    serialIterations:u32,
-    offsets:ptr<function,OffsetStruct>,
-    keys:KeyStruct
-){
-    if (gtid >= subgroup_size * serialIterations) {
-        let t = getWaveIndex(gtid) / serialIterations * HALF_RADIX;
-        
-        for (var i:u32 = 0; i < KEYS_PER_THREAD; i++)
-        {
-            let t2 = ExtractPackedIndex(keys.k[i]);
-            offsets.o[i] += ExtractPackedValue(atomicLoad(&g_d[t2 + t]) + atomicLoad(&g_d[t2]), keys.k[i]);
-        }
-    } else {
-        for (var i:u32 = 0; i < KEYS_PER_THREAD; i++){
-            offsets.o[i] += ExtractPackedValue(atomicLoad(&g_d[ExtractPackedIndex(keys.k[i])]), keys.k[i]);
-        }
-    }
-}
-
-fn ScatterKeysShared( offsets:OffsetStruct, keys:KeyStruct )
-{
-    for (var i:u32 = 0; i < KEYS_PER_THREAD; i++){
-        atomicStore(&g_d[offsets.o[i]], keys.k[i]);
-    }
-}
-
-fn DescendingIndex( deviceIndex:u32 ) -> u32
-{
-    return settings.e_numKeys - deviceIndex - 1;
-}
-
-fn WriteKey(deviceIndex:u32, groupSharedIndex:u32)
-{
-    b_alt[deviceIndex] = atomicLoad(&g_d[groupSharedIndex]);
-}
-
-fn LoadPayload( payload:ptr<function, u32>, deviceIndex:u32)
-{
-    *payload = b_sortPayload[deviceIndex];
-}
-
-fn ScatterPayloadsShared( offsets:OffsetStruct, payloads:KeyStruct)
-{
-    ScatterKeysShared(offsets, payloads);
-}
-
-fn WritePayload( deviceIndex:u32, groupSharedIndex:u32)
-{
-    b_altPayload[deviceIndex] = atomicLoad(&g_d[groupSharedIndex]);
-}
-
-//*****************************************************************************
-//SCATTERING: FULL PARTITIONS
-//*****************************************************************************
-
-fn ScatterPairsKeyPhaseAscending(
-    gtid:u32,
-    digits: ptr<function,DigitStruct>)
-{
-    var t = gtid;
-    for (var i:u32 = 0; i < KEYS_PER_THREAD; i++)
-    {
-        digits.d[i] = ExtractDigit(atomicLoad(&g_d[t]));
-        WriteKey(atomicLoad(&g_d[digits.d[i] + PART_SIZE]) + t, t);
-        t += D_DIM;
-    }
-}
-
-fn LoadPayloadsWGE16(
-    gtid:u32,
-    partIndex:u32,
-    payloads:ptr<function ,KeyStruct>
-){
-    var t = DeviceOffsetWGE16(gtid, partIndex);
-    for (var i:u32 = 0;
-        i < KEYS_PER_THREAD;
-        i++
-    ){
-        LoadPayload(&payloads.k[i], t);
-        t += subgroup_size;
-    }
-}
-
-fn LoadPayloadsWLT16(
-    gtid:u32,
-    partIndex:u32,
-    serialIterations:u32,
-    payloads:ptr<function ,KeyStruct>
-){
-    var t = DeviceOffsetWLT16(gtid, partIndex, serialIterations);
-    for (var i:u32 = 0;
-        i < KEYS_PER_THREAD;
-        i++
-    ){
-        LoadPayload(&payloads.k[i], t);
-        t += subgroup_size * serialIterations;
-    }
-}
-
-fn ScatterPayloadsAscending( gtid:u32 , digits: DigitStruct )
-{
-    var t = gtid;
-    for (var i:u32 = 0; i < KEYS_PER_THREAD; i++){
-        WritePayload(atomicLoad(&g_d[digits.d[i] + PART_SIZE]) + t, t);
-        t += D_DIM;
-    }
-}
-
-fn LoadPayloadsPartialWGE16(
-    gtid:u32,
-    partIndex:u32,
-    payloads:ptr<function,KeyStruct>
-){
-    var t = DeviceOffsetWGE16(gtid, partIndex);
-    for (var i:u32 = 0;
-        i < KEYS_PER_THREAD;
-        i++
-    ){
-        if (t < settings.e_numKeys){
-            LoadPayload(&payloads.k[i], t);
-        }
-        t += subgroup_size;
-    }
-}
-
-fn LoadPayloadsPartialWLT16(
-    gtid:u32,
-    partIndex:u32,
-    serialIterations:u32,
-    payloads: ptr<function,KeyStruct>
-){
-    var t = DeviceOffsetWLT16(gtid, partIndex, serialIterations);
-    for (var i:u32 = 0;
-        i < KEYS_PER_THREAD;
-        i++
-    ){
-        if (t < settings.e_numKeys){
-            LoadPayload(&payloads.k[i], t);
-        }
-        
-        t += subgroup_size * serialIterations;
-    }
-}
-
-fn ScatterPayloadsAscendingPartial(
-    gtid:u32,
-    finalPartSize:u32,
-    digits:DigitStruct
-){
-    var t = gtid;
-    for (var i:u32 = 0; i < KEYS_PER_THREAD; i++ )
-    {
-        
-        if (t < finalPartSize){
-            WritePayload(atomicLoad(&g_d[digits.d[i] + PART_SIZE]) + t, t);
-        }
-        
-        t += D_DIM;
-    }
-}
-
-
-fn ScatterDevicePartial(
-    gtid:u32,
-    partIndex:u32,
-    offsets: OffsetStruct
-){
-    var digits:DigitStruct;
-    let finalPartSize = settings.e_numKeys - partIndex * PART_SIZE;
-    ScatterPairsKeyPhaseAscendingPartial(gtid, finalPartSize, &digits);
-    workgroupBarrier();
-    
-    var payloads:KeyStruct;
-    if (subgroup_size >= 16){
-        LoadPayloadsPartialWGE16(gtid, partIndex, &payloads);
-    }else{
-        LoadPayloadsPartialWLT16(gtid, partIndex, SerialIterations(), &payloads);
-    }
-    ScatterPayloadsShared(offsets, payloads);
-    workgroupBarrier();
-    
-    ScatterPayloadsAscendingPartial(gtid, finalPartSize, digits);
-}
-//*****************************************************************************
-//DIGIT BINNING PASS KERNEL
-//*****************************************************************************
-fn AssignPartitionTile( gtid:u32 )-> u32
-{
-    if (gtid == 0){
-        atomicStore(&g_d[D_TOTAL_SMEM - 1], atomicAdd(&b_index[CurrentPass()], 1));
-    }
-    
-    workgroupBarrier();
-    
-    return atomicLoad(&g_d[D_TOTAL_SMEM - 1]);
-}
-
-fn DeviceBroadcastReductionsWGE16( gtid:u32,  partIndex:u32,  histReduction:u32)
-{
-    if (partIndex < settings.e_threadBlocks - 1)
-    {
-        atomicAdd(&b_passHist[gtid + PassHistOffset(partIndex + 1)],
-            (FLAG_REDUCTION | histReduction) << 2);
-    }
-}
-
-fn DeviceBroadcastReductionsWLT16( gtid:u32,  partIndex:u32,  histReduction:u32)
-{
-    if (partIndex < settings.e_threadBlocks - 1)
-    {
-        atomicAdd(&b_passHist[(gtid << 1) + PassHistOffset(partIndex + 1)],
-            (FLAG_REDUCTION | (histReduction & 0xffff)) << 2);
-                
-        atomicAdd(&b_passHist[(gtid << 1) + 1 + PassHistOffset(partIndex + 1)],
-            (FLAG_REDUCTION | ((histReduction >> 16) & 0xffff)) << 2);
-    }
-}
-
-
-fn Lookback(gtid:u32, partIndex:u32, exclusiveHistReduction:u32)
-{
-    if (gtid < RADIX)
-    {
-        var lookbackReduction:u32 = 0;
-        for (var k = partIndex; k >= 0;)
-        {
-            let flagPayload:u32 = atomicLoad(&b_passHist[gtid + PassHistOffset(k)]);
-            if ((flagPayload & FLAG_MASK) == FLAG_INCLUSIVE)
-            {
-                lookbackReduction += flagPayload >> 2;
-                if (partIndex < settings.e_threadBlocks - 1)
-                {
-                    atomicAdd(&b_passHist[gtid + PassHistOffset(partIndex + 1)],
-                        (1 | lookbackReduction) << 2);
-                }
-                atomicStore(&g_d[gtid + PART_SIZE], lookbackReduction - exclusiveHistReduction);
-                break;
-            }
-                    
-            if ((flagPayload & FLAG_MASK) == FLAG_REDUCTION)
-            {
-                lookbackReduction += flagPayload >> 2;
-                k--;
-            }
-        }
-    }
-}
-
-fn ScatterPairsKeyPhaseAscendingPartial(
-    gtid:u32,
-    finalPartSize:u32,
-    digits: ptr<function,DigitStruct>
-){
-    var t = gtid;
-    for (var i = 0; i < KEYS_PER_THREAD; i++)
-    {
-        if (t < finalPartSize)
-        {
-            let digit = ExtractDigit(atomicLoad(&g_d[t]));
-            digits.d[i] = digit;
-            WriteKey(atomicLoad(&g_d[digit + PART_SIZE])  + t, t);
-        }
-        t += D_DIM;
-    }
-}
-
-@compute @workgroup_size(N_DIM, 1, 1)
-fn main(
-    @builtin(local_invocation_index) gtid: u32,
-    @builtin(subgroup_invocation_id) subgroup_id:u32,
-    @builtin(subgroup_size) sg_size:u32,
-){
-    subgroup_index = subgroup_id;
-
-    // Propagate the subgroup-size builtin into the module-scope private var that every helper
-    // function reads. It was never written (the entry param used to shadow it), so all helpers
-    // divided by zero. Per-invocation store of an identical value — no shared memory, no barrier.
-    subgroup_size = sg_size;
-
-    // Uniformity: every condition or trip count that guards a workgroupBarrier must be provably
-    // workgroup-uniform. Reads of the var<private> subgroup_size are treated as non-uniform by
-    // the analysis, so anything barrier-adjacent is derived HERE from the (uniform) builtin and
-    // passed down / used directly.
-    let serial_iterations = (D_DIM / sg_size + 31) >> 5;
-    let wave_hists_size_wge16 = D_DIM / sg_size * RADIX;
-
-    var partitionIndex: u32;
-    var keys: KeyStruct;
-    var offsets: OffsetStruct;
-
-    //WGT 16 can potentially skip some barriers
-    if (sg_size > 16)
-    {
-        if(wave_hists_size_wge16 < PART_SIZE){
-            ClearWaveHists(gtid);
-        }
-
-        partitionIndex = AssignPartitionTile(gtid);
-        if(wave_hists_size_wge16 >= PART_SIZE)
-        {
-            workgroupBarrier();
-            ClearWaveHists(gtid);
-            workgroupBarrier();
-        }
-    }
-    
-    if (sg_size <= 16)
-    {
-        partitionIndex = AssignPartitionTile(gtid);
-        workgroupBarrier();
-        ClearWaveHists(gtid);
-        workgroupBarrier();
-    }
-
-    if(partitionIndex < settings.e_threadBlocks - 1)
-    {
-        if (sg_size >= 16){
-            keys = LoadKeysWGE16(gtid, partitionIndex);
-        }
-
-        if (sg_size < 16){
-            keys = LoadKeysWLT16(gtid, partitionIndex, serial_iterations);
-        }
-    }
-
-    if(partitionIndex == settings.e_threadBlocks - 1)
-    {
-        if (sg_size >= 16){
-            keys = LoadKeysPartialWGE16(gtid, partitionIndex);
-        }
-
-        if (sg_size < 16){
-            keys = LoadKeysPartialWLT16(gtid, partitionIndex, serial_iterations);
-        }
-    }
-
-    var exclusiveHistReduction:u32;
-    if (sg_size >= 16)
-    {
-        offsets = RankKeysWGE16(gtid, keys);
-        workgroupBarrier();
-        
-        var histReduction:u32;
-        if (gtid < RADIX)
-        {
-            histReduction = WaveHistInclusiveScanCircularShiftWGE16(gtid);
-            DeviceBroadcastReductionsWGE16(gtid, partitionIndex, histReduction);
-            
-        }
-        
-        let a01 = subgroupExclusiveAdd(histReduction);
-        
-        if (gtid < RADIX){
-            histReduction += a01; //take advantage of barrier to begin scan
-        }
-        
-        
-        workgroupBarrier();
-
-        WaveHistReductionExclusiveScanWGE16(gtid, histReduction);
-        workgroupBarrier();
-            
-        UpdateOffsetsWGE16(gtid, &offsets, keys);
-        if (gtid < RADIX){
-            exclusiveHistReduction = atomicLoad(&g_d[gtid]); //take advantage of barrier to grab value
-        }
-        workgroupBarrier();
-    }
-    
-    if (sg_size < 16)
-    {
-        offsets = RankKeysWLT16(gtid, keys, serial_iterations);
-
-        if (gtid < HALF_RADIX)
-        {
-            let histReduction = WaveHistInclusiveScanCircularShiftWLT16(gtid);
-            atomicStore(&g_d[gtid], histReduction + (histReduction << 16)); //take advantage of barrier to begin scan
-            DeviceBroadcastReductionsWLT16(gtid, partitionIndex, histReduction);
-        }
-
-        WaveHistReductionExclusiveScanWLT16(gtid);
-        workgroupBarrier();
-
-        UpdateOffsetsWLT16(gtid, serial_iterations, &offsets, keys);
-        if (gtid < RADIX){ //take advantage of barrier to grab value
-            exclusiveHistReduction = (atomicLoad(&g_d[gtid >> 1]) >> select(0u, 16u, (gtid & 1) != 0 ) ) & 0xffff;
-        }
-        workgroupBarrier();
-    }
-    
-    ScatterKeysShared(offsets, keys);
-    Lookback(gtid, partitionIndex, exclusiveHistReduction);
-    workgroupBarrier();
-    
-    let last_block_index = settings.e_threadBlocks - 1;
-    {
-        // ==========
-        // unrolls ScatterDevice to achieve flow uniformity (WebGPU requirement)
-        // ==========
-        var digits:DigitStruct;
-        if(partitionIndex < last_block_index){
-            ScatterPairsKeyPhaseAscending(gtid, &digits);
-        }
-        
-        workgroupBarrier();
-            
-        var payloads:KeyStruct;
-        
-        if(partitionIndex < last_block_index){
-            if (sg_size >= 16){
-                LoadPayloadsWGE16(gtid, partitionIndex, &payloads);
-            }else{
-                LoadPayloadsWLT16(gtid, partitionIndex, serial_iterations, &payloads);
-            }
-            ScatterPayloadsShared(offsets, payloads);
-        }
-            
-            workgroupBarrier();
-        
-        if(partitionIndex < last_block_index){    
-            ScatterPayloadsAscending(gtid, digits);
-        }
-    }
-    
-    { 
-        // ==========
-        // unrolls ScatterDevicePartial to achieve flow uniformity (WebGPU requirement)
-        // ==========
-           
-        var digits:DigitStruct;
-        let finalPartSize = settings.e_numKeys - partitionIndex * PART_SIZE;
-        
-        let is_last_partition = partitionIndex == last_block_index;
-        
-        if(is_last_partition){
-            ScatterPairsKeyPhaseAscendingPartial(gtid, finalPartSize, &digits);
-        }
-        
-        workgroupBarrier();
-        
-        if(is_last_partition){
-            var payloads:KeyStruct;
-            if (sg_size >= 16){
-                LoadPayloadsPartialWGE16(gtid, partitionIndex, &payloads);
-            }else{
-                LoadPayloadsPartialWLT16(gtid, partitionIndex, serial_iterations, &payloads);
-            }
-            ScatterPayloadsShared(offsets, payloads);
-        }
-        
-        workgroupBarrier();
-            
-        if(is_last_partition){    
-            ScatterPayloadsAscendingPartial(gtid, finalPartSize, digits);
-        }
-        
-    }
-}
-    `
-);
-
-body.addExtension(WebGPUExtensionType.Subgroups);
-
-export const shader_digit_binning_pass_pairs = ComputeShader.from({
-    label: "RadixSort / Digit Binning Pass",
-    resources,
-    body,
-});