submission 500137

"""
Prefix Sum - DLB with Warp-Parallel Lookback
@MemoryCoalesced

v16: Eliminate boundary-check ISETPs from hot path.

PTX analysis showed 147 ISETP instructions — almost all from:
  if (gf + 3 < n) ... else scalar fallback
on every load and store. For n=268,435,456 = 16384 * TILE_SIZE,
every tile is full so these checks are dead code in the contest.

Fix: template on bool FULL_TILE.
  - FULL_TILE=true:  pure float4 loads/stores, zero ISETPs, zero branches
  - FULL_TILE=false: original scalar fallback for tail tile
Dispatch: if n % TILE_SIZE == 0, all tiles take FULL_TILE=true path.
"""

from task import input_t, output_t
import torch
from torch.utils.cpp_extension import load_inline

cpp_source = """
torch::Tensor prefix_sum_cuda(torch::Tensor input, torch::Tensor output);
"""

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define BLOCK_THREADS    256
#define ITEMS_PER_THREAD 64
#define TILE_SIZE        (BLOCK_THREADS * ITEMS_PER_THREAD)  // 16384

#define STATUS_INVALID   0
#define STATUS_PARTIAL   1
#define STATUS_INCLUSIVE 2

__device__ __forceinline__ long long pack(int status, float value) {
    return ((long long)(unsigned)status << 32) | (long long)__float_as_uint(value);
}
__device__ __forceinline__ void unpack(long long packed, int &status, float &value) {
    status = (int)(unsigned int)(packed >> 32);
    value  = __uint_as_float((unsigned int)(packed & 0xFFFFFFFFULL));
}

__device__ __forceinline__ float warp_inclusive_scan(float val) {
    #pragma unroll
    for (int offset = 1; offset < 32; offset <<= 1) {
        float tmp = __shfl_up_sync(0xffffffff, val, offset);
        if ((threadIdx.x & 31) >= offset) val += tmp;
    }
    return val;
}

__device__ float block_inclusive_scan(float val, float* smem_scan) {
    const int lane   = threadIdx.x & 31;
    const int warp   = threadIdx.x >> 5;
    const int nwarps = BLOCK_THREADS / 32;

    val = warp_inclusive_scan(val);
    if (lane == 31) smem_scan[warp] = val;
    __syncthreads();

    if (warp == 0) {
        float w = (lane < nwarps) ? smem_scan[lane] : 0.0f;
        w = warp_inclusive_scan(w);
        if (lane < nwarps) smem_scan[lane] = w;
    }
    __syncthreads();

    if (warp > 0) val += smem_scan[warp - 1];
    return val;
}

__device__ __forceinline__ float warp_reduce_add(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xffffffff, val, offset);
    return val;
}

__device__ float warp_parallel_lookback(
    volatile long long* __restrict__ descriptors,
    int tile_id)
{
    const int lane = threadIdx.x & 31;
    float prefix = 0.0f;
    int base = tile_id - 1;

    while (true) {
        int look_idx = base - lane;
        int   status = STATUS_INCLUSIVE;
        float value  = 0.0f;

        if (look_idx >= 0) {
            long long desc;
            do {
                desc = descriptors[look_idx];
                unpack(desc, status, value);
            } while (status == STATUS_INVALID);
        }

        unsigned inclusive_mask = __ballot_sync(0xffffffff, status == STATUS_INCLUSIVE);
        if (inclusive_mask) {
            int first_inc = __ffs(inclusive_mask) - 1;
            float contrib = (lane <= first_inc) ? value : 0.0f;
            prefix += warp_reduce_add(contrib);
            return prefix;
        }

        prefix += warp_reduce_add(value);
        base -= 32;
        if (base < 0) break;
    }
    return prefix;
}

template <bool FULL_TILE>
__global__ void __launch_bounds__(BLOCK_THREADS)
scan_kernel(
    const float* __restrict__ input,
    float*       __restrict__ output,
    volatile long long* __restrict__ descriptors,
    int n)
{
    __shared__ float smem_scan[8];
    __shared__ float s_aggregate;
    __shared__ float s_prefix;

    const int tile_id      = blockIdx.x;
    const int tile_start   = tile_id * TILE_SIZE;
    const int thread_start = tile_start + threadIdx.x * ITEMS_PER_THREAD;
    constexpr int VEC = ITEMS_PER_THREAD / 4;

    // ----------------------------------------------------------------
    // 1. Load — float4 throughout, boundary checks only in tail tile
    // ----------------------------------------------------------------
    float items[ITEMS_PER_THREAD];
    const float4* in4 = reinterpret_cast<const float4*>(input + thread_start);

    #pragma unroll
    for (int v = 0; v < VEC; v++) {
        float4 val;
        if constexpr (FULL_TILE) {
            // No branch, no ISETP — compiler emits pure LDG.E.128
            val = __ldg(in4 + v);
        } else {
            int gf = thread_start + v * 4;
            if (gf + 3 < n) {
                val = __ldg(in4 + v);
            } else {
                val.x = (gf     < n) ? __ldg(input + gf    ) : 0.0f;
                val.y = (gf + 1 < n) ? __ldg(input + gf + 1) : 0.0f;
                val.z = (gf + 2 < n) ? __ldg(input + gf + 2) : 0.0f;
                val.w = (gf + 3 < n) ? __ldg(input + gf + 3) : 0.0f;
            }
        }
        items[v*4+0] = val.x; items[v*4+1] = val.y;
        items[v*4+2] = val.z; items[v*4+3] = val.w;
    }

    // ----------------------------------------------------------------
    // 2. Thread-local scan
    // ----------------------------------------------------------------
    #pragma unroll
    for (int i = 1; i < ITEMS_PER_THREAD; i++) items[i] += items[i-1];
    float thread_sum = items[ITEMS_PER_THREAD - 1];

    // ----------------------------------------------------------------
    // 3. Block scan
    // ----------------------------------------------------------------
    float block_result = block_inclusive_scan(thread_sum, smem_scan);

    // ----------------------------------------------------------------
    // 4. Publish descriptor
    // ----------------------------------------------------------------
    if (threadIdx.x == BLOCK_THREADS - 1) {
        s_aggregate = block_result;
        long long packed = (tile_id == 0)
            ? pack(STATUS_INCLUSIVE, block_result)
            : pack(STATUS_PARTIAL,   block_result);
        descriptors[tile_id] = packed;
        __threadfence();
    }
    __syncthreads();

    // ----------------------------------------------------------------
    // 5. Warp 0: lookback + promote
    // ----------------------------------------------------------------
    if (tile_id > 0 && (threadIdx.x >> 5) == 0) {
        float lbprefix = warp_parallel_lookback(descriptors, tile_id);
        if (threadIdx.x == 0) {
            s_prefix = lbprefix;
            descriptors[tile_id] = pack(STATUS_INCLUSIVE, lbprefix + s_aggregate);
            __threadfence();
        }
    } else if (tile_id == 0 && threadIdx.x == 0) {
        s_prefix = 0.0f;
    }
    __syncthreads();

    // ----------------------------------------------------------------
    // 6. Apply prefix + store
    // ----------------------------------------------------------------
    float my_prefix = s_prefix + (block_result - thread_sum);
    #pragma unroll
    for (int i = 0; i < ITEMS_PER_THREAD; i++) items[i] += my_prefix;

    float4* out4 = reinterpret_cast<float4*>(output + thread_start);
    #pragma unroll
    for (int v = 0; v < VEC; v++) {
        float4 val = { items[v*4+0], items[v*4+1], items[v*4+2], items[v*4+3] };
        if constexpr (FULL_TILE) {
            out4[v] = val;
        } else {
            int gf = thread_start + v * 4;
            if (gf + 3 < n) {
                out4[v] = val;
            } else {
                if (gf     < n) output[gf    ] = val.x;
                if (gf + 1 < n) output[gf + 1] = val.y;
                if (gf + 2 < n) output[gf + 2] = val.z;
                if (gf + 3 < n) output[gf + 3] = val.w;
            }
        }
    }
}

// Explicit instantiations so both are compiled
template __global__ void scan_kernel<true>(const float*, float*, volatile long long*, int);
template __global__ void scan_kernel<false>(const float*, float*, volatile long long*, int);

torch::Tensor prefix_sum_cuda(torch::Tensor input, torch::Tensor output) {
    const int n         = input.numel();
    const int num_tiles = (n + TILE_SIZE - 1) / TILE_SIZE;

    auto desc = torch::zeros({num_tiles},
        torch::TensorOptions().dtype(torch::kInt64).device(input.device()));

    auto* d_in  = input.data_ptr<float>();
    auto* d_out = output.data_ptr<float>();
    auto* d_desc = (volatile long long*)desc.data_ptr<int64_t>();

    if (n % TILE_SIZE == 0) {
        // All tiles full — zero boundary checks
        scan_kernel<true><<<num_tiles, BLOCK_THREADS>>>(d_in, d_out, d_desc, n);
    } else {
        // Last tile is partial — use safe path for all tiles
        // (could split: full path for [0, num_tiles-1], safe for last — but
        //  this case doesn't occur in the contest so not worth the complexity)
        scan_kernel<false><<<num_tiles, BLOCK_THREADS>>>(d_in, d_out, d_desc, n);
    }

    return output;
}
"""

module = None

def get_module():
    global module
    if module is None:
        module = load_inline(
            name='dlb_scan_v16',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['prefix_sum_cuda'],
            verbose=False,
            extra_cuda_cflags=[
                '-O3',
                '--use_fast_math',
                '-std=c++17',
            ],
        )
    return module

def custom_kernel(data: input_t) -> output_t:
    inp, out = data
    mod = get_module()
    mod.prefix_sum_cuda(inp, out)
    return out