submission 507438

"""
Sort v47 — Fused single-kernel + full_tile fast path + alt-bits

Single kernel: histogram + prefix scan + persistent scatter all fused.
256 threads × IPT=52 = TILE 13312 | 7513 tiles | 2 blocks/SM
Full tile fast path eliminates bounds checks on 7,512 of 7,513 tiles.
Alternating status bits eliminate cooperative lookback zeroing between passes.
Bigger tiles = fewer lookback chains = less global memory polling.

1.879ms on H100
"""

import torch
from torch.utils.cpp_extension import load_inline

cuda_source = r"""
#include <cuda_runtime.h>
#include <stdint.h>

#define RADIX        256
#define MAX_PASSES   4
#define THREADS      256
#define IPT          52
#define TILE         (THREADS * IPT)
#define WARPS        (THREADS / 32)
#define WARP_SEG     (IPT * 32)
#define DYN_SMEM     (TILE * 4)

#define MAX_N        100000000
#define MAX_NB       ((MAX_N + TILE - 1) / TILE)

#define STATUS_MASK     0xC0000000u
#define STATUS_PARTIAL  0x40000000u
#define VAL_MASK        0x3FFFFFFFu

__device__ __forceinline__ uint32_t lb_load(const uint32_t* addr) {
    uint32_t v;
    asm volatile("ld.relaxed.gpu.global.b32 %0, [%1];" : "=r"(v) : "l"(addr));
    return v;
}

__device__ __forceinline__ uint32_t float_to_uint(uint32_t u) {
    return u ^ ((u >> 31) ? 0xFFFFFFFFu : 0x80000000u);
}
__device__ __forceinline__ float uint_to_float(uint32_t u) {
    u ^= ((u >> 31) ? 0x80000000u : 0xFFFFFFFFu);
    return __uint_as_float(u);
}


__global__ void __launch_bounds__(THREADS, 2)
sort_fused(
    const float*    __restrict__ inp_float,
    float*          __restrict__ out_float,
    uint32_t*       __restrict__ buf0,
    uint32_t*       __restrict__ buf1,
    uint32_t*       __restrict__ global_hist,
    uint32_t*       __restrict__ global_offsets,
    uint32_t*       __restrict__ lookback,
    uint32_t*       __restrict__ counters,
    uint32_t*       __restrict__ barrier,
    uint32_t*       __restrict__ hist_barrier,
    int nb, int n, int num_blocks)
{
    extern __shared__ uint32_t s_exchange[];

    __shared__ uint32_t s_warp_hist[WARPS * RADIX];
    __shared__ uint32_t s_warp_scan[WARPS * RADIX];
    __shared__ uint32_t s_count[RADIX];
    __shared__ uint32_t s_prefix[RADIX];
    __shared__ uint32_t s_tile_prefix[RADIX];
    __shared__ uint32_t s_scan_tmp[WARPS];
    __shared__ uint32_t s_tile_id;

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int bid  = blockIdx.x;

    // PHASE A: Cooperative global histogram
    {
        uint32_t* s_hist = s_exchange;
        for (int i = tid; i < MAX_PASSES * RADIX; i += THREADS)
            s_hist[i] = 0u;
        __syncthreads();

        for (int tile = bid; tile < nb; tile += num_blocks) {
            const int tbase = tile * TILE;
            const int tsz = min(TILE, n - tbase);
            const uint4* src4 = reinterpret_cast<const uint4*>(inp_float + tbase);

            #pragma unroll
            for (int i = 0; i < IPT / 4; i++) {
                int base_f = (i * THREADS + tid) * 4;
                if (base_f + 3 < tsz) {
                    uint4 v = src4[i * THREADS + tid];
                    uint32_t kx=float_to_uint(v.x),ky=float_to_uint(v.y);
                    uint32_t kz=float_to_uint(v.z),kw=float_to_uint(v.w);
                    atomicAdd(&s_hist[0*RADIX+((kx>> 0)&0xFF)],1u); atomicAdd(&s_hist[1*RADIX+((kx>> 8)&0xFF)],1u);
                    atomicAdd(&s_hist[2*RADIX+((kx>>16)&0xFF)],1u); atomicAdd(&s_hist[3*RADIX+((kx>>24)&0xFF)],1u);
                    atomicAdd(&s_hist[0*RADIX+((ky>> 0)&0xFF)],1u); atomicAdd(&s_hist[1*RADIX+((ky>> 8)&0xFF)],1u);
                    atomicAdd(&s_hist[2*RADIX+((ky>>16)&0xFF)],1u); atomicAdd(&s_hist[3*RADIX+((ky>>24)&0xFF)],1u);
                    atomicAdd(&s_hist[0*RADIX+((kz>> 0)&0xFF)],1u); atomicAdd(&s_hist[1*RADIX+((kz>> 8)&0xFF)],1u);
                    atomicAdd(&s_hist[2*RADIX+((kz>>16)&0xFF)],1u); atomicAdd(&s_hist[3*RADIX+((kz>>24)&0xFF)],1u);
                    atomicAdd(&s_hist[0*RADIX+((kw>> 0)&0xFF)],1u); atomicAdd(&s_hist[1*RADIX+((kw>> 8)&0xFF)],1u);
                    atomicAdd(&s_hist[2*RADIX+((kw>>16)&0xFF)],1u); atomicAdd(&s_hist[3*RADIX+((kw>>24)&0xFF)],1u);
                } else {
                    for (int fi = 0; fi < 4; fi++) {
                        if (base_f + fi < tsz) {
                            uint32_t k = float_to_uint(__float_as_uint(inp_float[tbase + base_f + fi]));
                            atomicAdd(&s_hist[0*RADIX+((k>> 0)&0xFF)],1u); atomicAdd(&s_hist[1*RADIX+((k>> 8)&0xFF)],1u);
                            atomicAdd(&s_hist[2*RADIX+((k>>16)&0xFF)],1u); atomicAdd(&s_hist[3*RADIX+((k>>24)&0xFF)],1u);
                        }
                    }
                }
            }
        }
        __syncthreads();
        for (int i = tid; i < MAX_PASSES * RADIX; i += THREADS)
            if (s_hist[i] > 0u) atomicAdd(&global_hist[i], s_hist[i]);
    }

    // PHASE B: Prefix scan
    __threadfence();
    if (tid == 0) {
        atomicAdd(hist_barrier, 1u);
        while (lb_load(hist_barrier) < (uint32_t)num_blocks) {}
    }
    __syncthreads();

    if (bid == 0) {
        for (int pass = 0; pass < MAX_PASSES; pass++) {
            s_count[tid] = global_hist[pass * RADIX + tid];
            __syncthreads();
            for (int stride = 1; stride < RADIX; stride <<= 1) {
                uint32_t v = (tid >= stride) ? s_count[tid - stride] : 0u;
                __syncthreads(); s_count[tid] += v; __syncthreads();
            }
            global_offsets[pass * RADIX + tid] = (tid > 0) ? s_count[tid - 1] : 0u;
            __syncthreads();
        }
        __threadfence();
        if (tid == 0) atomicAdd(hist_barrier, 1u);
    }
    if (tid == 0) {
        while (lb_load(hist_barrier) < (uint32_t)num_blocks + 1u) {}
    }
    __syncthreads();

    // PHASE C: Persistent scatter with full_tile fast path
    const int shifts[4] = {0, 8, 16, 24};

    for (int pass = 0; pass < 4; pass++) {
        const int shift = shifts[pass];
        const uint32_t* src;
        uint32_t*       dst;
        const bool is_first = (pass == 0);
        const bool is_last  = (pass == 3);

        const uint32_t STATUS_INVALID  = (uint32_t)(pass & 1) << 31;
        const uint32_t STATUS_COMPLETE = STATUS_INVALID ^ 0x80000000u;

        if (pass == 0) { src = (const uint32_t*)inp_float; dst = buf0; }
        else if (pass == 1) { src = buf0; dst = buf1; }
        else if (pass == 2) { src = buf1; dst = buf0; }
        else { src = buf0; dst = (uint32_t*)out_float; }

        const uint32_t* pass_go = global_offsets + pass * RADIX;
        uint32_t* pass_lb = lookback;

        while (true) {
            if (tid == 0) s_tile_id = atomicAdd(&counters[pass], 1u);
            __syncthreads();
            const uint32_t tile_id = s_tile_id;
            if (tile_id >= (uint32_t)nb) break;

            const int tbase = (int)(tile_id * TILE);
            const int wbase = tbase + warp * WARP_SEG;
            const bool full_tile = (tbase + TILE <= n);

            for (int i = tid; i < WARPS * RADIX; i += THREADS) {
                s_warp_hist[i] = 0u;
                s_warp_scan[i] = 0u;
            }
            __syncthreads();

            uint32_t keys[IPT];
            uint32_t ranks[IPT];

            if (full_tile) {
                #pragma unroll
                for (int i = 0; i < IPT; i++) {
                    int gidx = wbase + i * 32 + lane;
                    if (is_first)
                        keys[i] = float_to_uint(__float_as_uint(((const float*)src)[gidx]));
                    else
                        keys[i] = src[gidx];
                }

                #pragma unroll
                for (int i = 0; i < IPT; i++) {
                    uint32_t d = (keys[i] >> shift) & 0xFF;
                    ranks[i] = atomicAdd(&s_warp_hist[warp * RADIX + d], 1u);
                }
            } else {
                #pragma unroll
                for (int i = 0; i < IPT; i++) {
                    int gidx = wbase + i * 32 + lane;
                    if (gidx < n) {
                        if (is_first)
                            keys[i] = float_to_uint(__float_as_uint(((const float*)src)[gidx]));
                        else
                            keys[i] = src[gidx];
                    } else {
                        keys[i] = 0u;
                    }
                }

                #pragma unroll
                for (int i = 0; i < IPT; i++) {
                    int gidx = wbase + i * 32 + lane;
                    ranks[i] = 0u;
                    if (gidx < n) {
                        uint32_t d = (keys[i] >> shift) & 0xFF;
                        ranks[i] = atomicAdd(&s_warp_hist[warp * RADIX + d], 1u);
                    }
                }
            }
            __syncthreads();

            if (tid < RADIX) {
                uint32_t run = 0u;
                #pragma unroll
                for (int w = 0; w < WARPS; w++) {
                    s_warp_scan[w * RADIX + tid] = run;
                    run += s_warp_hist[w * RADIX + tid];
                }
                s_count[tid] = run;
            }
            __syncthreads();

            if (tid < RADIX) {
                uint32_t my_count = s_count[tid];
                uint32_t exclusive = 0u;

                if (tile_id == 0u) {
                    atomicExch(&pass_lb[0u * RADIX + tid], STATUS_COMPLETE | my_count);
                } else {
                    atomicExch(&pass_lb[tile_id * RADIX + tid], STATUS_PARTIAL | my_count);
                    __threadfence();
                    int look = (int)tile_id - 1;
                    while (look >= 0) {
                        uint32_t v;
                        do { v = lb_load(&pass_lb[look * RADIX + tid]); } while ((v & STATUS_MASK) == STATUS_INVALID);
                        exclusive += v & VAL_MASK;
                        if ((v & STATUS_MASK) != STATUS_PARTIAL) break;
                        look--;
                    }
                    atomicExch(&pass_lb[tile_id * RADIX + tid],
                               STATUS_COMPLETE | (exclusive + my_count));
                }
                s_prefix[tid] = pass_go[tid] + exclusive;
            }
            __syncthreads();

            {
                uint32_t val = s_count[tid];
                #pragma unroll
                for (int off = 1; off < 32; off <<= 1) {
                    uint32_t n2 = __shfl_up_sync(0xFFFFFFFF, val, off);
                    if (lane >= off) val += n2;
                }
                if (lane == 31) s_scan_tmp[warp] = val;
            }
            __syncthreads();
            if (tid == 0) {
                uint32_t run = 0u;
                for (int w = 0; w < WARPS; w++) {
                    uint32_t t = s_scan_tmp[w]; s_scan_tmp[w] = run; run += t;
                }
            }
            __syncthreads();
            {
                uint32_t val = s_count[tid];
                #pragma unroll
                for (int off = 1; off < 32; off <<= 1) {
                    uint32_t n2 = __shfl_up_sync(0xFFFFFFFF, val, off);
                    if (lane >= off) val += n2;
                }
                s_tile_prefix[tid] = val - s_count[tid] + s_scan_tmp[warp];
            }
            __syncthreads();

            if (full_tile) {
                #pragma unroll
                for (int i = 0; i < IPT; i++) {
                    uint32_t d = (keys[i] >> shift) & 0xFF;
                    uint32_t rank = s_warp_scan[warp * RADIX + d] + ranks[i];
                    uint32_t smem_pos = s_tile_prefix[d] + rank;
                    s_exchange[smem_pos ^ (smem_pos >> 5)] = keys[i];
                }
            } else {
                #pragma unroll
                for (int i = 0; i < IPT; i++) {
                    int gidx = wbase + i * 32 + lane;
                    if (gidx < n) {
                        uint32_t d = (keys[i] >> shift) & 0xFF;
                        uint32_t rank = s_warp_scan[warp * RADIX + d] + ranks[i];
                        uint32_t smem_pos = s_tile_prefix[d] + rank;
                        s_exchange[smem_pos ^ (smem_pos >> 5)] = keys[i];
                    }
                }
            }
            __syncthreads();

            if (full_tile) {
                #pragma unroll 4
                for (int i = 0; i < IPT; i++) {
                    int p = i * THREADS + tid;
                    uint32_t k = s_exchange[p ^ (p >> 5)];
                    uint32_t d = (k >> shift) & 0xFF;
                    uint32_t out_pos = s_prefix[d] + ((uint32_t)p - s_tile_prefix[d]);
                    if (is_last)
                        ((float*)dst)[out_pos] = uint_to_float(k);
                    else
                        dst[out_pos] = k;
                }
            } else {
                const int tile_count = min((int)TILE, n - tbase);
                #pragma unroll 4
                for (int i = 0; i < IPT; i++) {
                    int p = i * THREADS + tid;
                    if (p < tile_count) {
                        uint32_t k = s_exchange[p ^ (p >> 5)];
                        uint32_t d = (k >> shift) & 0xFF;
                        uint32_t out_pos = s_prefix[d] + ((uint32_t)p - s_tile_prefix[d]);
                        if (is_last)
                            ((float*)dst)[out_pos] = uint_to_float(k);
                        else
                            dst[out_pos] = k;
                    }
                }
            }
            __syncthreads();
        }

        __threadfence();
        if (tid == 0) {
            atomicAdd(&barrier[pass], 1u);
            while (lb_load(&barrier[pass]) < (uint32_t)num_blocks) {}
        }
        __syncthreads();
    }
}

#include <torch/extension.h>

static bool g_configured = false;
static int g_num_persistent = 0;
static void configure_once() {
    cudaFuncSetAttribute(sort_fused, cudaFuncAttributeMaxDynamicSharedMemorySize, DYN_SMEM);
    int blocks_per_sm = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &blocks_per_sm, sort_fused, THREADS, DYN_SMEM);
    int num_sms = 0;
    cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0);
    g_num_persistent = num_sms * blocks_per_sm;
    g_configured = true;
}

void sort_cuda_py(
    torch::Tensor inp,
    torch::Tensor out,
    torch::Tensor tmp,
    torch::Tensor global_hist,
    torch::Tensor global_offsets,
    torch::Tensor lookback,
    torch::Tensor counter,
    torch::Tensor barrier,
    torch::Tensor hist_barrier)
{
    if (!g_configured) configure_once();

    const int n  = inp.numel();
    const int nb = (n + TILE - 1) / TILE;
    const int num_persistent = g_num_persistent;

    float*    d_inp = inp.data_ptr<float>();
    float*    d_out = out.data_ptr<float>();
    uint32_t* d_buf0 = (uint32_t*)tmp.data_ptr<int>();
    uint32_t* d_buf1 = reinterpret_cast<uint32_t*>(out.data_ptr<float>());
    uint32_t* d_gh  = (uint32_t*)global_hist.data_ptr<int>();
    uint32_t* d_go  = (uint32_t*)global_offsets.data_ptr<int>();
    uint32_t* d_lb  = (uint32_t*)lookback.data_ptr<int>();
    uint32_t* d_ctr = (uint32_t*)counter.data_ptr<int>();
    uint32_t* d_bar = (uint32_t*)barrier.data_ptr<int>();
    uint32_t* d_hbar = (uint32_t*)hist_barrier.data_ptr<int>();

    cudaMemsetAsync(d_gh,   0, MAX_PASSES * RADIX * sizeof(uint32_t));
    cudaMemsetAsync(d_ctr,  0, 4 * sizeof(uint32_t));
    cudaMemsetAsync(d_bar,  0, 4 * sizeof(uint32_t));
    cudaMemsetAsync(d_hbar, 0, sizeof(uint32_t));
    cudaMemsetAsync(d_lb,   0, (size_t)nb * RADIX * sizeof(uint32_t));

    sort_fused<<<num_persistent, THREADS, DYN_SMEM>>>(
        d_inp, d_out,
        d_buf0, d_buf1,
        d_gh, d_go,
        d_lb, d_ctr, d_bar, d_hbar,
        nb, n, num_persistent);
}
"""

cpp_source = """
void sort_cuda_py(
    torch::Tensor inp, torch::Tensor out, torch::Tensor tmp,
    torch::Tensor global_hist, torch::Tensor global_offsets,
    torch::Tensor lookback, torch::Tensor counter,
    torch::Tensor barrier, torch::Tensor hist_barrier);
"""

_MAX_N  = 100_000_000
_TILE   = 256 * 52
_MAX_NB = (_MAX_N + _TILE - 1) // _TILE

_tmp            = torch.empty(_MAX_N,            device='cuda', dtype=torch.int32)
_global_hist    = torch.zeros(4 * 256,           device='cuda', dtype=torch.int32)
_global_offsets = torch.zeros(4 * 256,           device='cuda', dtype=torch.int32)
_lookback       = torch.zeros(_MAX_NB * 256,     device='cuda', dtype=torch.int32)
_counter        = torch.zeros(4,                 device='cuda', dtype=torch.int32)
_barrier        = torch.zeros(4,                 device='cuda', dtype=torch.int32)
_hist_barrier   = torch.zeros(1,                 device='cuda', dtype=torch.int32)

module = load_inline(
    name='sort_v47_submit',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['sort_cuda_py'],
    verbose=True,
    extra_cuda_cflags=['-O3', '--use_fast_math', '-std=c++17', '--ptxas-options=-v'],
)

def custom_kernel(data: tuple) -> torch.Tensor:
    inp, out = data
    module.sort_cuda_py(
        inp, out,
        _tmp[:inp.numel()],
        _global_hist,
        _global_offsets,
        _lookback,
        _counter,
        _barrier,
        _hist_barrier,
    )
    return out