submission 799493

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200

# Blocked compact-WY Householder QR, batched, with SHARED-MEMORY-RESIDENT
# panel factorization. The bottleneck in the global-memory version was the
# within-panel rank-1 updates hammering gmem; here each CTA loads its
# (n-c) x b panel into smem (B200 opt-in 200KB), factors entirely in smem,
# builds the WY T, writes back. Trailing update = batched fp32 GEMM (bmm).
# Block width b chosen per n so the panel fits smem. Large-n/small-batch and
# tiny-n route to torch.geqrf.
import os
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

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

__device__ __forceinline__ float warpRed(float s) {
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) s += __shfl_down_sync(0xFFFFFFFFu, s, o);
    return s;
}

// smem layout: [P: rows*b][vs: rows][S: b*b][red:32]
__global__ void bpanel_smem(float* __restrict__ Aall, float* __restrict__ Tall,
                            float* __restrict__ tauAll, int n, int c, int b) {
    const int m = blockIdx.x;
    float* __restrict__ A = Aall + (size_t)m * n * n;
    float* __restrict__ T = Tall + (size_t)m * b * b;
    float* __restrict__ tau = tauAll + (size_t)m * n;
    const int t = threadIdx.x, NT = blockDim.x, lane = t & 31, warp = t >> 5;
    const int rows = n - c;

    extern __shared__ float sm[];
    float* P = sm;                 // rows x b, row-major P[i*b+jl]
    float* vs = P + (size_t)rows * b;
    float* S = vs + rows;          // b x b
    float* red = S + (size_t)b * b;
    __shared__ float s_beta, s_tau, s_inv;

    for (int idx = t; idx < rows * b; idx += NT) {
        int i = idx / b, jl = idx % b;
        P[idx] = A[(size_t)(c + i) * n + (c + jl)];
    }
    __syncthreads();

    for (int jl = 0; jl < b; jl++) {
        float s = 0.f;
        for (int i = jl + t; i < rows; i += NT) { float a = P[i * b + jl]; s += a * a; }
        s = warpRed(s);
        if (lane == 0) red[warp] = s;
        __syncthreads();
        if (t == 0) {
            float nsq = 0.f; int nw = (NT + 31) >> 5;
            for (int w = 0; w < nw; w++) nsq += red[w];
            float alpha = P[jl * b + jl];
            float tailsq = nsq - alpha * alpha;
            if (tailsq <= 0.f) { s_tau = 0.f; s_beta = alpha; s_inv = 0.f; }
            else { float xn = sqrtf(nsq); float beta = (alpha >= 0.f) ? -xn : xn;
                   s_beta = beta; s_tau = (beta - alpha) / beta; s_inv = 1.f / (alpha - beta); }
        }
        __syncthreads();
        const float tauj = s_tau, beta = s_beta, inv = s_inv;
        if (t == 0) vs[jl] = 1.f;
        for (int i = jl + 1 + t; i < rows; i += NT) vs[i] = P[i * b + jl] * inv;
        __syncthreads();
        if (tauj != 0.f) {
            for (int k = jl + 1 + t; k < b; k += NT) {
                float w = 0.f;
                for (int i = jl; i < rows; i++) w += vs[i] * P[i * b + k];
                w *= tauj;
                for (int i = jl; i < rows; i++) P[i * b + k] -= w * vs[i];
            }
        }
        __syncthreads();
        if (t == 0) { P[jl * b + jl] = beta; tau[c + jl] = tauj; }
        for (int i = jl + 1 + t; i < rows; i += NT) P[i * b + jl] = vs[i];
        __syncthreads();
    }

    // write panel back
    for (int idx = t; idx < rows * b; idx += NT) {
        int i = idx / b, jl = idx % b;
        A[(size_t)(c + i) * n + (c + jl)] = P[idx];
    }

    // ---- build WY T (b x b upper-tri) from smem P ----
    for (int idx = t; idx < b * b; idx += NT) S[idx] = 0.f;
    __syncthreads();
    for (int pair = warp; pair < b * b; pair += (NT >> 5)) {
        int il = pair / b, jl = pair % b;
        if (il >= jl) continue;
        // S[il,jl] = sum_{i>=jl} V[i,il]*V[i,jl];  V[jl,jl]=1, V[i,il]=P[i*b+il]
        float acc = 0.f;
        // i == jl term: P[jl*b+il] * 1
        if (lane == 0) acc += P[jl * b + il];
        for (int i = jl + 1 + lane; i < rows; i += 32) acc += P[i * b + il] * P[i * b + jl];
        acc = warpRed(acc);
        if (lane == 0) S[il * b + jl] = acc;
    }
    __syncthreads();
    if (t == 0) {
        for (int jl = 0; jl < b; jl++) {
            float tj = tau[c + jl];
            float z[128];
            for (int il = 0; il < jl; il++) z[il] = -tj * S[il * b + jl];
            for (int il = 0; il < jl; il++) {
                float acc = 0.f;
                for (int kl = il; kl < jl; kl++) acc += T[il * b + kl] * z[kl];
                T[il * b + jl] = acc;
            }
            T[jl * b + jl] = tj;
        }
    }
}

static int g_smem_set = 0;
void bpanel_launch(torch::Tensor A, torch::Tensor T, torch::Tensor tau,
                   int64_t n, int64_t c, int64_t b) {
    int B = A.size(0);
    int rows = n - c;
    int threads = 512;
    size_t smem = ((size_t)rows * b + rows + (size_t)b * b + 32) * sizeof(float);
    if (!g_smem_set) {
        cudaFuncSetAttribute(bpanel_smem, cudaFuncAttributeMaxDynamicSharedMemorySize, 200 * 1024);
        g_smem_set = 1;
    }
    bpanel_smem<<<B, threads, smem>>>(A.data_ptr<float>(), T.data_ptr<float>(),
                                      tau.data_ptr<float>(), (int)n, (int)c, (int)b);
}
"""

_CPP = "void bpanel_launch(torch::Tensor, torch::Tensor, torch::Tensor, int64_t, int64_t, int64_t);"

_mod = load_inline(
    name="qr_bpanel_smem",
    cpp_sources=_CPP, cuda_sources=_CUDA,
    functions=["bpanel_launch"],
    extra_cuda_cflags=["-O3", "-arch=sm_100a", "--use_fast_math"],
    extra_cflags=["-O3"], verbose=False)

_LO = int(os.environ.get("QR_LO", "129"))
_HI = int(os.environ.get("QR_HI", "1536"))


def _bw(n):
    # panel width so that rows*b*4 fits ~180KB at the first (largest) panel
    if n <= 1024:
        return 32
    return 32


def _blocked(A, n):
    B = A.shape[0]
    b = _bw(n)
    tau = torch.zeros(B, n, device=A.device, dtype=torch.float32)
    eye = torch.eye(b, device=A.device, dtype=torch.float32)
    for c in range(0, n, b):
        bb = min(b, n - c)
        Tm = torch.zeros(B, bb, bb, device=A.device, dtype=torch.float32)
        _mod.bpanel_launch(A, Tm, tau, n, c, bb)
        ce = c + bb
        if ce < n:
            Vblk = A[:, c:n, c:ce]
            Vtop = Vblk[:, :bb, :].tril(-1) + eye[:bb, :bb]
            V = torch.cat([Vtop, Vblk[:, bb:, :]], dim=1)
            At = A[:, c:n, ce:n]
            Y = torch.bmm(V.transpose(1, 2), At)
            Z = torch.bmm(Tm.triu().transpose(1, 2), Y)
            At.add_(torch.bmm(V, Z), alpha=-1.0)
    return tau


def custom_kernel(data: input_t) -> output_t:
    n = data.shape[-1]
    if _LO <= n <= _HI:
        A = data.contiguous().clone()
        tau = _blocked(A, n)
        return (A, tau)
    return torch.geqrf(data)