submission 808354

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200

# qr_v2: batched square Householder QR matching torch.geqrf(data) -> (H, tau).
#
# Iteration 6: blocked Householder, general panel width PB via compact-WY T matrix.
#   - One CTA per matrix, whole batch in one launch (custom path n<=512).
#   - Process columns in panels of PB (=4). Panel factorization is unblocked but
#     restricted to the panel; then build the PB x PB upper-triangular T (LARFT,
#     forward/columnwise) so that H_j0 ... H_{j0+PB-1} = I - V T V^T, and apply the
#     block reflector to the trailing matrix as:
#         C := (I - V T^T V^T) C  =>  w = V^T c ; g = T^T w ; c -= V g
#     This applies PB reflectors per 2 passes over the trailing matrix (vs PB
#     separate rank-1 passes) -> ~PB/2 less trailing-matrix traffic & barriers than
#     b=2, attacking the latency/critical-path bound (kernel ~100x below FP32 peak).
#   - V panel staged in shared ROW-MAJOR (Vbuf[r*PB + c]) so the inner loop over c
#     in the trailing update hits consecutive banks (no conflicts). geqrf-exact
#     packing. n>512 falls back to torch.geqrf.
#
# VARIANT panelfuse: fuse the Phase-1 within-panel reflector applications into ONE
#   multi-column block reduction. After computing reflector jj, instead of a runtime
#   loop "for k=j+1..j0+pb-1 { blockReduceSum }", each thread accumulates partial dot
#   products for ALL remaining panel columns into a register array part[PB] using a
#   FULLY #pragma unroll loop over the COMPILE-TIME constant PB with PREDICATION of
#   inactive columns. A single multi-value block reduction (warp-shuffle each of the
#   PB partials -> shared scratch[nwarps*PB] -> combine in warp 0, 2 __syncthreads)
#   yields all dots; then one fused update pass, again compile-time-unrolled with
#   predication. No per-thread array is ever indexed by a runtime bound (no local
#   spill). Phase 2/3 unchanged.

import torch
from torch.utils.cpp_extension import load_inline

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

#define MAXW 32  // max warps per block (1024 threads)
// PB (panel width) is a TEMPLATE parameter on blockReduceSumPB / qr_kernel, dispatched
// per-n by the launcher (iter-30): PB=4 for small n (<=256), PB=8 for large n. Measured
// crossover: small matrices want narrow PB (n=32 -33%, n=176 -21% at PB=4), large want wide.

// Block-wide sum reduction using warp shuffles + shared scratch (<=32 warps).
__device__ __forceinline__ float blockReduceSum(float val, float* sm){
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) val += __shfl_down_sync(0xffffffffu, val, o);
    if (lane == 0) sm[wid] = val;
    __syncthreads();
    int nwarps = (blockDim.x + 31) >> 5;
    val = (threadIdx.x < nwarps) ? sm[lane] : 0.0f;
    if (wid == 0) {
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) val += __shfl_down_sync(0xffffffffu, val, o);
    }
    if (threadIdx.x == 0) sm[0] = val;
    __syncthreads();
    float r = sm[0];
    // NOTE: no trailing protect-sync. The only caller (Phase-1 norm) always executes a
    // __syncthreads (line ~132) between this sm[0] read and the next sm_red write, so no
    // thread can clobber sm[0] before every thread has read it. Saves 1 barrier/column.
    return r;
}

// Multi-value block reduction: reduce PB partials at once into out[PB].
// scratch must have room for nwarps*PB floats. Costs 2 __syncthreads.
template<int PB>
__device__ __forceinline__ void blockReduceSumPB(float part[PB], float* scratch, float out[PB]){
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    int nwarps = (blockDim.x + 31) >> 5;
    // warp-shuffle each partial
    #pragma unroll
    for (int c = 0; c < PB; ++c) {
        float v = part[c];
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
        if (lane == 0) scratch[wid * PB + c] = v;
    }
    __syncthreads();
    // warp 0 combines per-warp partials
    if (wid == 0) {
        #pragma unroll
        for (int c = 0; c < PB; ++c) {
            float acc = 0.0f;
            for (int w = lane; w < nwarps; w += 32) acc += scratch[w * PB + c];
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, o);
            // safe: writing scratch[c] (==scratch[0*PB+c]) only after all reads
            // of column c are done; later iterations read column c'!=c.
            if (lane == 0) scratch[c] = acc;
        }
    }
    __syncthreads();
    #pragma unroll
    for (int c = 0; c < PB; ++c) out[c] = scratch[c];
    // NOTE: no trailing protect-sync. Every caller (Phase-1 apply line ~169; Phase-2 lines
    // ~202/213) executes a __syncthreads between this scratch read and the next reduction's
    // scratch write, so scratch cannot be clobbered before all threads read it.
}

// One block factors one (n x n) matrix in row-major storage (A[i][j] = Ab[i*n+j]).
// Blocked compact-WY Householder, panel width PB.
template<int PB>
__global__ void qr_kernel(float* __restrict__ A, float* __restrict__ tau, int n){
    extern __shared__ float sh[];          // Vbuf[r*PB+c] : PB*n ; then T[i*PB+c] : PB*PB
    float* Vbuf = sh;
    float* Tsh  = sh + (size_t)PB * n;
    __shared__ float sm_red[32];
    __shared__ float sm_redPB[MAXW * PB];
    int bidx = blockIdx.x;
    float* Ab   = A   + (size_t)bidx * n * n;
    float* taub = tau + (size_t)bidx * n;
    int tid = threadIdx.x;
    int nt  = blockDim.x;

    for (int j0 = 0; j0 < n; j0 += PB) {
        int pb = min(PB, n - j0);

        // ===== Phase 1: panel factorization (unblocked, restricted to panel) =====
        for (int jj = 0; jj < pb; ++jj) {
            int j = j0 + jj;
            float alpha = Ab[(size_t)j * n + j];
            float p = 0.0f;
            for (int i = j + 1 + tid; i < n; i += nt) { float a = Ab[(size_t)i * n + j]; p += a * a; }
            float xs = blockReduceSum(p, sm_red);
            float tauj, beta, inv;
            if (xs == 0.0f) { tauj = 0.0f; beta = alpha; inv = 0.0f; }
            else { float r = sqrtf(alpha * alpha + xs); beta = (alpha >= 0.0f) ? -r : r;
                   tauj = (beta - alpha) / beta; inv = 1.0f / (alpha - beta); }
            // store reflector jj into Vbuf column jj: 0 above row j (within [j0,j)), 1 at j, tail below
            for (int i = j0 + tid; i < j; i += nt) Vbuf[(size_t)i * PB + jj] = 0.0f;
            if (xs != 0.0f) {
                for (int i = j + 1 + tid; i < n; i += nt) { float vi = Ab[(size_t)i * n + j] * inv; Ab[(size_t)i * n + j] = vi; Vbuf[(size_t)i * PB + jj] = vi; }
            } else {
                for (int i = j + 1 + tid; i < n; i += nt) Vbuf[(size_t)i * PB + jj] = 0.0f;
            }
            if (tid == 0) { Vbuf[(size_t)j * PB + jj] = 1.0f; Ab[(size_t)j * n + j] = beta; taub[j] = tauj; }
            __syncthreads();

            // ---- FUSED apply of reflector jj to remaining panel columns ----
            // Active columns are k = j+1 .. j0+pb-1, i.e. local offset c=0..(pb-jj-2).
            // Predicate: column c (=> k = j+1+c) is active iff (j+1+c) < (j0+pb).
            // Accumulate partial dots for ALL PB columns; index part[] only with
            // compile-time-constant c. Rows i in [j, n) (Vbuf is 0 for i in [j+1.. )? no:
            // reflector v has v[j]=1 and tail below; rows above j are not part of this
            // reflector. Use the same row range as base: i = j + tid .. n-1.
            int nrem = j0 + pb - (j + 1);   // number of active remaining columns (>=0)
            if (nrem > 0) {
                float part[PB];
                #pragma unroll
                for (int c = 0; c < PB; ++c) part[c] = 0.0f;
                for (int i = j + tid; i < n; i += nt) {
                    float vi = Vbuf[(size_t)i * PB + jj];
                    const float* arow = Ab + (size_t)i * n + (j + 1);
                    #pragma unroll
                    for (int c = 0; c < PB; ++c) {
                        // predicate: only accumulate for active columns
                        float aval = (c < nrem) ? arow[c] : 0.0f;
                        part[c] += vi * aval;
                    }
                }
                float wsh[PB];
                blockReduceSumPB<PB>(part, sm_redPB, wsh);
                #pragma unroll
                for (int c = 0; c < PB; ++c) wsh[c] *= tauj;
                // fused update pass
                for (int i = j + tid; i < n; i += nt) {
                    float vi = Vbuf[(size_t)i * PB + jj];
                    float* arow = Ab + (size_t)i * n + (j + 1);
                    #pragma unroll
                    for (int c = 0; c < PB; ++c) {
                        if (c < nrem) arow[c] -= wsh[c] * vi;
                    }
                }
                __syncthreads();
            }
        }

        // ===== Phase 2: build T (pb x pb upper-triangular), LARFT forward/columnwise =====
        if (tid < pb) Tsh[(size_t)tid * PB + tid] = taub[j0 + tid];   // diagonal = tau
        __syncthreads();
        for (int c = 1; c < pb; ++c) {
            float tauc = taub[j0 + c];
            // tmp[i] = -tauc * (V[:,i]^T V[:,c]) for i=0..c-1, stored into Tsh[i*PB+c].
            // FUSED (panelfuse-style): compute ALL c dot products in ONE pass over rows +
            // ONE multi-value block reduction, instead of c separate blockReduceSum calls
            // (was pb(pb-1)/2 reductions/panel ~ 3 syncs each; now pb-1 ~ 2 syncs each).
            // V[:,c] is loaded once per row (was c times); part[] indexed only by the
            // compile-time-constant i with predication of inactive (i>=c) columns -> no spill.
            float part[PB];
            #pragma unroll
            for (int i = 0; i < PB; ++i) part[i] = 0.0f;
            for (int r = j0 + c + tid; r < n; r += nt) {
                float vc = Vbuf[(size_t)r * PB + c];
                const float* vrow = Vbuf + (size_t)r * PB;
                #pragma unroll
                for (int i = 0; i < PB; ++i) {
                    float vi = (i < c) ? vrow[i] : 0.0f;
                    part[i] += vi * vc;
                }
            }
            float yv[PB];
            blockReduceSumPB<PB>(part, sm_redPB, yv);   // ends with __syncthreads
            if (tid == 0) {
                #pragma unroll
                for (int i = 0; i < PB; ++i) if (i < c) Tsh[(size_t)i * PB + c] = -tauc * yv[i];
            }
            __syncthreads();
            // T[0:c,c] = T[0:c,0:c] @ tmp  (upper-tri matvec) ; do serially on thread 0
            if (tid == 0) {
                float col[PB];
                for (int i = 0; i < c; ++i) {
                    float acc = 0.0f;
                    for (int l = i; l < c; ++l) acc += Tsh[(size_t)i * PB + l] * Tsh[(size_t)l * PB + c];
                    col[i] = acc;
                }
                for (int i = 0; i < c; ++i) Tsh[(size_t)i * PB + c] = col[i];
            }
            __syncthreads();
        }

        // ===== Phase 3: trailing update  C -= V (T^T (V^T C))  for k = j0+pb .. n-1 =====
        for (int k = j0 + pb + tid; k < n; k += nt) {
            float w[PB];
            #pragma unroll
            for (int c = 0; c < PB; ++c) w[c] = 0.0f;
            // Pass 1: w = V^T C[:,k]. One thread owns column k and walks all rows; the
            // Ab[r*n+k] loads are strided (no cross-thread reuse) and latency-bound, so
            // unroll by 4 to keep 4 independent global loads in flight (memory-level
            // parallelism) instead of stalling on each load before the next.
            int r = j0;
            float* Ck = Ab + k;                  // &Ab[0*n + k]; row r at Ck + r*n (pass 2 writes it)
            for (; r + 4 <= n; r += 4) {
                float c0 = Ck[(size_t)(r + 0) * n];
                float c1 = Ck[(size_t)(r + 1) * n];
                float c2 = Ck[(size_t)(r + 2) * n];
                float c3 = Ck[(size_t)(r + 3) * n];
                const float* v0 = Vbuf + (size_t)(r + 0) * PB;
                const float* v1 = Vbuf + (size_t)(r + 1) * PB;
                const float* v2 = Vbuf + (size_t)(r + 2) * PB;
                const float* v3 = Vbuf + (size_t)(r + 3) * PB;
                #pragma unroll
                for (int c = 0; c < PB; ++c)
                    w[c] += v0[c] * c0 + v1[c] * c1 + v2[c] * c2 + v3[c] * c3;
            }
            for (; r < n; ++r) {
                float cval = Ck[(size_t)r * n];
                #pragma unroll
                for (int c = 0; c < PB; ++c) w[c] += Vbuf[(size_t)r * PB + c] * cval;
            }
            float g[PB];
            for (int c = 0; c < pb; ++c) {
                float acc = 0.0f;
                for (int i = 0; i <= c; ++i) acc += Tsh[(size_t)i * PB + c] * w[i];
                g[c] = acc;
            }
            // Pass 2: C[:,k] -= V g. Same strided read-modify-write; unroll by 4 so the
            // 4 reads issue together before the dependent FMAs + stores.
            r = j0;
            for (; r + 4 <= n; r += 4) {
                float* p0 = Ck + (size_t)(r + 0) * n;
                float* p1 = Ck + (size_t)(r + 1) * n;
                float* p2 = Ck + (size_t)(r + 2) * n;
                float* p3 = Ck + (size_t)(r + 3) * n;
                float a0 = *p0, a1 = *p1, a2 = *p2, a3 = *p3;
                const float* v0 = Vbuf + (size_t)(r + 0) * PB;
                const float* v1 = Vbuf + (size_t)(r + 1) * PB;
                const float* v2 = Vbuf + (size_t)(r + 2) * PB;
                const float* v3 = Vbuf + (size_t)(r + 3) * PB;
                float u0 = 0.0f, u1 = 0.0f, u2 = 0.0f, u3 = 0.0f;
                #pragma unroll
                for (int c = 0; c < pb; ++c) {
                    float gc = g[c];
                    u0 += gc * v0[c]; u1 += gc * v1[c]; u2 += gc * v2[c]; u3 += gc * v3[c];
                }
                *p0 = a0 - u0; *p1 = a1 - u1; *p2 = a2 - u2; *p3 = a3 - u3;
            }
            for (; r < n; ++r) {
                float upd = 0.0f;
                for (int c = 0; c < pb; ++c) upd += g[c] * Vbuf[(size_t)r * PB + c];
                Ck[(size_t)r * n] -= upd;
            }
        }
        __syncthreads();
    }
}

void qr_batched(torch::Tensor H, torch::Tensor tau){
    int batch = H.size(0);
    int n     = H.size(1);
    // Block size selection (unchanged): n=1024 -> 1024 threads (1 trailing col/thread);
    // n<=512 -> threads = round_up(n,32) so tiny matrices don't pay idle-warp sync cost.
    int threads = (n > 512) ? 1024 : (((n + 31) / 32) * 32);
    // Per-n panel width (iter-30 dispatch). Small matrices have a tiny trailing region, so
    // per-panel overhead dominates (Phase-2 T-build is O(PB^2)/panel; Phase-1 within-panel
    // apply is O(PB)/reflector) -> a NARROW panel wins. Large matrices are trailing-traffic-
    // bound (total C re-traffic ~ n^3/PB) -> a WIDE panel wins. Measured crossover sits
    // between n=176 (PB=4: n=32 -33%, n=176 -21%) and n=352 (PB=8). EXACTLY ONE kernel
    // instance launches per call (the branches are mutually exclusive).
    if (n <= 256) {
        const int PBv = 4;
        size_t shmem = ((size_t)PBv * n + (size_t)PBv * PBv) * sizeof(float);
        qr_kernel<PBv><<<batch, threads, shmem>>>(H.data_ptr<float>(), tau.data_ptr<float>(), n);
    } else {
        const int PBv = 8;
        size_t shmem = ((size_t)PBv * n + (size_t)PBv * PBv) * sizeof(float);
        qr_kernel<PBv><<<batch, threads, shmem>>>(H.data_ptr<float>(), tau.data_ptr<float>(), n);
    }
}
'''

_cpp_src = r"void qr_batched(torch::Tensor H, torch::Tensor tau);"

_mod = load_inline(
    name="qr_v2_pbdispatch",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["qr_batched"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


def custom_kernel(data):
    # data: (batch, n, n) float32 CUDA tensor. Return (H, tau) per torch.geqrf.
    if data.dim() != 3 or data.shape[-1] != data.shape[-2] or data.shape[-1] > 1024:
        return torch.geqrf(data)
    batch, n, _ = data.shape
    H = data.contiguous().clone()
    tau = torch.empty((batch, n), dtype=torch.float32, device=data.device)
    _mod.qr_batched(H, tau)
    return H, tau