submission 822973

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200

# Batched square compact-Householder QR (geqrf-compatible (H, tau) output).
#
# Two execution paths, routed on shape only (never on matrix values, so the
# "mixed" anti-cheat rule is satisfied) -- both run the same backward-stable
# Householder math, so every conditioning profile passes the FP32 gate:
#
#   * Blocked batched (n <= 2048): right-looking blocked Householder with the WY
#     representation. A fused panel-factorization kernel (one block per matrix,
#     panel resident in shared memory) produces R, the reflectors V, tau and the
#     pb x pb WY factor T in-kernel; the trailing update is three batched cuBLAS
#     SGEMMs (BLAS-3). Keeping the T-factor in-kernel matters on the leaderboard
#     runtime, where per-launch overhead is high -- the blocked path's launch
#     count, not its FLOPs, is what dominates there, so we minimize launches
#     (everything the panel can do, it does in one kernel).
#
#   * cuSOLVER (n > 2048, i.e. n=4096 b=2): too few, too large matrices -- the
#     blocked panel would under-occupy the GPU and explode the per-panel launch
#     count, so each matrix goes through cuSOLVER's own blocked geqrf in a loop.
#     (cuSOLVER is column-major, so this path clones transposed and returns a
#     transposed view.)
#
# Everything runs in the default execution queue the harness provides; no extra
# queues are created.

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

CUDA_SRC = r"""
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>
#include <cusolverDn.h>
#include <vector>

// ---------------------------------------------------------------------------
__device__ __forceinline__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
    return v;
}
__device__ __forceinline__ float block_reduce_sum(float val, float* scratch) {
    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
    val = warp_reduce_sum(val);
    if (lane == 0) scratch[wid] = val;
    __syncthreads();
    int nwarps = (blockDim.x + 31) >> 5;
    val = (threadIdx.x < nwarps) ? scratch[lane] : 0.0f;
    if (wid == 0) {
        val = warp_reduce_sum(val);
        if (lane == 0) scratch[0] = val;
    }
    __syncthreads();
    float t = scratch[0];
    __syncthreads();
    return t;
}

// Fused panel factorization for a column-major matrix M (lda = n).
// Factors the panel rows[j, n) x cols[j, j+pb): writes R (panel upper incl diag),
// essential reflectors (panel strictly-lower), tau[j..j+pb), and the WY T-factor
// (pb x pb, upper-triangular) into Tg (stride pb*pb per matrix, column-major pb).
// One block per matrix; panel resident in shared memory.
//
// Shared layout (floats): P[h*pb] | BT[pb*pb] | scr[<=32] | (h <= n)
__global__ void panel_factor(float* __restrict__ A, float* __restrict__ tau,
                             float* __restrict__ Tg, float* __restrict__ Vg,
                             int n, int j, int pb) {
    const int b   = blockIdx.x;
    const int tid = threadIdx.x;
    const int nt  = blockDim.x;
    const int h   = n - j;                 // panel height
    float* __restrict__ M = A + (long long)b * n * n;
    float* __restrict__ T = Tg + (long long)b * pb * pb;
    float* __restrict__ tA = tau + (long long)b * n;

    extern __shared__ float smem[];
    float* P    = smem;                    // h*pb, column-major (P[r + c*h])
    float* BT   = P + (long long)h * pb;   // pb*pb (B then reused for T), col-major
    float* zcol = BT + pb * pb;            // pb (LARFT temp)
    float* scr  = zcol + pb;               // <=32

    // 1. load panel into shared. M is ROW-MAJOR: element (row j+r, col j+c) at
    //    (j+r)*n + (j+c). Shared P stays column-major (P[r + c*h]).
    for (long long idx = tid; idx < (long long)h * pb; idx += nt) {
        int c = idx / h, r = idx - (long long)c * h;
        P[r + (long long)c * h] = M[(long long)(j + r) * n + (j + c)];
    }
    __syncthreads();

    // 2. factor the panel in shared (BLAS-2)
    for (int k = 0; k < pb; ++k) {
        // reflector for column k, rows [k, h)
        float local = 0.0f;
        for (int r = k + 1 + tid; r < h; r += nt) {
            float x = P[r + (long long)k * h];
            local += x * x;
        }
        float xnorm2 = block_reduce_sum(local, scr);
        float alpha = P[k + (long long)k * h];
        __syncthreads();
        float tauk;
        if (xnorm2 > 0.0f) {
            float beta = -copysignf(sqrtf(alpha * alpha + xnorm2), alpha);
            tauk = (beta - alpha) / beta;
            float inv = 1.0f / (alpha - beta);
            if (tid == 0) { P[k + (long long)k * h] = beta; tA[j + k] = tauk; }
            for (int r = k + 1 + tid; r < h; r += nt)
                P[r + (long long)k * h] *= inv;
        } else {
            tauk = 0.0f;
            if (tid == 0) tA[j + k] = 0.0f;
        }
        __syncthreads();
        // apply H_k = I - tauk v v^T (v[k]=1, v[r]=P[r,k]) to columns k+1..pb-1
        if (tauk != 0.0f) {
            int wid = tid >> 5, lane = tid & 31, nwarps = nt >> 5;
            for (int c = k + 1 + wid; c < pb; c += nwarps) {
                float w = 0.0f;
                // r=k term: v=1
                if (lane == 0) w += P[k + (long long)c * h];
                for (int r = k + 1 + lane; r < h; r += 32)
                    w += P[r + (long long)k * h] * P[r + (long long)c * h];
                w = warp_reduce_sum(w);
                w = __shfl_sync(0xffffffffu, w, 0) * tauk;
                if (lane == 0) P[k + (long long)c * h] -= w;          // r=k, v=1
                for (int r = k + 1 + lane; r < h; r += 32)
                    P[r + (long long)c * h] -= w * P[r + (long long)k * h];
            }
        }
        __syncthreads();
    }

    // 3. write panel back to M (row-major) and materialize V (unit lower-trapezoidal)
    //    into Vg row-major (h x pb, row stride pb, matrix stride n*pb).
    float* __restrict__ V = Vg + (long long)b * n * pb;
    for (long long idx = tid; idx < (long long)h * pb; idx += nt) {
        int c = idx / h, r = idx - (long long)c * h;
        float val = P[r + (long long)c * h];
        M[(long long)(j + r) * n + (j + c)] = val;
        V[(long long)r * pb + c] = (r < c) ? 0.0f : (r == c ? 1.0f : val);
    }

    // 4. build B[m,c] = v_m^T v_c (m<c) into BT (col-major pb), upper part only.
    //    v_c has unit diag at row c: v_c[c]=1, v_c[r]=P[r,c] for r>c, 0 for r<c.
    //    B[m,c] = P[c,m]*1 + sum_{r=c+1..h-1} P[r,m]*P[r,c]   (since m<c)
    {
        int wid = tid >> 5, lane = tid & 31, nwarps = nt >> 5;
        int npair = pb * pb;
        for (int p = wid; p < npair; p += nwarps) {
            int c = p / pb, m = p - c * pb;
            if (m < c) {
                float s = 0.0f;
                if (lane == 0) s += P[c + (long long)m * h];     // r=c term
                for (int r = c + 1 + lane; r < h; r += 32)
                    s += P[r + (long long)m * h] * P[r + (long long)c * h];
                s = warp_reduce_sum(s);
                if (lane == 0) BT[m + c * pb] = s;
            }
        }
    }
    __syncthreads();

    // 5. LARFT forward recurrence -> T (overwrite BT). T upper-tri, col-major pb.
    //    z[p] = -tau_c * B[p,c]; T[0:c,c] = T[0:c,0:c] @ z; T[c,c]=tau_c.
    //    Sequential over c; threads parallel over rows m. zcol decouples the read
    //    of B[:,c] from the write of T[:,c] so there is no in-place race.
    for (int c = 0; c < pb; ++c) {
        float tauc = tA[j + c];
        for (int p = tid; p < c; p += nt)
            zcol[p] = -tauc * BT[p + c * pb];          // -tau_c * B[p,c]
        __syncthreads();
        for (int m = tid; m < c; m += nt) {
            float acc = 0.0f;
            for (int p = m; p < c; ++p)                // T[m,p] (p<c, built) * z[p]
                acc += BT[m + p * pb] * zcol[p];
            BT[m + c * pb] = acc;
        }
        __syncthreads();
        if (tid == 0) BT[c + c * pb] = tauc;
        __syncthreads();
    }

    // write T to global
    for (long long idx = tid; idx < (long long)pb * pb; idx += nt) {
        int c = idx / pb, m = idx - (long long)c * pb;
        T[m + c * pb] = (m <= c) ? BT[m + c * pb] : 0.0f;
    }
}

// ---------------------------------------------------------------------------
// Own library handles. A freshly created handle uses the default execution
// queue (the same one our <<<>>> kernels use) -- so kernels and library calls
// stay correctly ordered without ever naming or creating an extra queue.
static cusolverDnHandle_t g_solv = nullptr;
static cublasHandle_t     g_blas = nullptr;

#define CK_CUDA(x)   do { cudaError_t e_=(x); TORCH_CHECK(e_==cudaSuccess, "cuda: ", cudaGetErrorString(e_)); } while(0)
#define CK_BLAS(x)   do { cublasStatus_t s_=(x); TORCH_CHECK(s_==CUBLAS_STATUS_SUCCESS, "cublas status ", (int)s_); } while(0)
#define CK_SOLV(x)   do { cusolverStatus_t s_=(x); TORCH_CHECK(s_==CUSOLVER_STATUS_SUCCESS, "cusolver status ", (int)s_); } while(0)

std::vector<torch::Tensor> qr_compact(torch::Tensor A) {
    TORCH_CHECK(A.is_cuda() && A.dim() == 3 && A.size(1) == A.size(2));
    TORCH_CHECK(A.scalar_type() == torch::kFloat32);
    int batch = A.size(0), n = A.size(1);
    auto tau = torch::empty({batch, n}, A.options());
    float* Tp = tau.data_ptr<float>();
    if (batch == 0 || n == 0)
        return {A.clone(), tau};

    // Nothing below names or creates an execution queue: kernels launch in the
    // default queue and the cuBLAS / cuSOLVER handles default to it too, so every
    // kernel and library call stays correctly ordered.
    //
    // Route only the few-huge-matrix case (n>2048, i.e. n=4096 b=2) to cuSOLVER;
    // n=2048 b=8 goes through the blocked path (cuSOLVER's B200 FP32 geqrf is ~3x
    // slower there). cuSOLVER is column-major, so its path clones into a transposed
    // buffer and returns a transposed view; the blocked path below works row-major.
    bool use_cusolver = (n > 2048);
    if (use_cusolver) {
        auto At = A.transpose(-1, -2).contiguous();   // col-major working buffer
        float* Ap = At.data_ptr<float>();
        if (!g_solv) CK_SOLV(cusolverDnCreate(&g_solv));
        int lwork = 0;
        CK_SOLV(cusolverDnSgeqrf_bufferSize(g_solv, n, n, Ap, n, &lwork));
        auto work = torch::empty({lwork}, A.options());
        auto info = torch::empty({1}, torch::TensorOptions().dtype(torch::kInt32).device(A.device()));
        for (int b = 0; b < batch; ++b)
            CK_SOLV(cusolverDnSgeqrf(g_solv, n, n, Ap + (long long)b * n * n, n,
                             Tp + (long long)b * n, work.data_ptr<float>(), lwork,
                             info.data_ptr<int>()));
        CK_CUDA(cudaGetLastError());
        return {At.transpose(-1, -2), tau};
    }

    // Blocked batched path (row-major throughout; clone is a plain copy).
    auto At = A.clone();
    float* Ap = At.data_ptr<float>();
    int dev = A.device().index();
    int shmem_optin = at::cuda::getDeviceProperties(dev)->sharedMemPerBlockOptin;
    // Use the device's full opt-in shared memory (B200 = 227 KB). This lets the
    // panel fit for larger n (e.g. n=2048 -> pb=24), which is what makes the
    // blocked path beat cuSOLVER's slow B200 geqrf on the n=2048 batch.
    int budget = shmem_optin - 256;                // small margin (need[] has +64 pad)
    // pick pb (multiple of 8, capped at 48) so n*pb + pb*pb + pb + 32 floats fit.
    // Larger pb => fewer panels => fewer kernel/GEMM launches and a larger k in the
    // trailing GEMMs (better arithmetic intensity); the leaderboard runtime rewards
    // both. Capped at 48 to keep >=2 panel blocks resident per SM at n=512.
    int pb = 0;
    for (int cand = 8; cand <= 48; cand += 8) {
        long long need = ((long long)n * cand + cand * cand + cand + 64) * sizeof(float);
        if (need <= budget) pb = cand; else break;
    }
    if (pb == 0) pb = 8;
    if (pb > n) pb = n;

    auto Vg = torch::empty({batch, (long long)n * pb}, A.options());
    auto Tg = torch::empty({batch, (long long)pb * pb}, A.options());
    auto W1 = torch::empty({batch, (long long)pb * n}, A.options());
    auto W2 = torch::empty({batch, (long long)pb * n}, A.options());
    float* Vp = Vg.data_ptr<float>();
    float* Tgp = Tg.data_ptr<float>();
    float* W1p = W1.data_ptr<float>();
    float* W2p = W2.data_ptr<float>();

    // Use our OWN cuBLAS handle (default queue), not torch's getCurrentCUDABlasHandle
    // -- torch's handle is bound to torch's current queue, which the harness may set
    // to a non-default one for timing; that would run the GEMMs out of order with our
    // default-queue kernels. Our handle defaults to the same queue as the kernels.
    if (!g_blas) {
        CK_BLAS(cublasCreate(&g_blas));
        // Force true IEEE FP32 in the trailing-update GEMMs. On data-center GPUs
        // cuBLAS may otherwise use TF32 tensor cores for SGEMM (~10-bit mantissa),
        // which blows the FP32 residual gate on ill-conditioned / mixed batches.
        CK_BLAS(cublasSetMathMode(g_blas, CUBLAS_PEDANTIC_MATH));
    }
    cublasHandle_t hbl = g_blas;
    const float one = 1.0f, zero = 0.0f, neg = -1.0f;

    // Thread count by how well the batch fills the GPU. With many matrices
    // (batch >= 128) there are already enough blocks to saturate the 148 SMs, so
    // 256 threads/block is best and 512 only adds sync overhead (measured: n=512
    // b=640 regresses with 512). With few matrices (n=1024 b=60, n=2048 b=8) the
    // panel is block-starved, so 512 threads/block adds useful parallelism
    // (measured: n=1024 -3%, n=2048 -5%). Small n stays 256 (a wide block idles).
    int threads = (batch < 128 && n >= 256) ? 512 : 256;
    size_t max_shmem = ((long long)n * pb + pb * pb + pb + 64) * sizeof(float);
    CK_CUDA(cudaFuncSetAttribute(panel_factor, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)max_shmem));

    for (int j = 0; j < n; j += pb) {
        int pbj = (pb < n - j) ? pb : (n - j);
        int h = n - j;
        size_t shmem = ((long long)h * pbj + pbj * pbj + pbj + 64) * sizeof(float);
        panel_factor<<<batch, threads, shmem>>>(Ap, Tp, Tgp, Vp, n, j, pbj);
        CK_CUDA(cudaGetLastError());

        int jb = j + pbj;
        int tcols = n - jb;
        if (tcols > 0) {
            // Row-major trailing update C = C - V (T^T (V^T C)), C = M[j:n, jb:n].
            // cuBLAS is column-major; the col-major view of row-major X (lda=row
            // stride) is X^T, so we operate on transposes:
            //   Ccm=C^T (tcols x h, lda=n), Vcm=V^T (pb x h, lda=pb), Tcm-view=T.
            //   W1cm = Ccm * V         (=W1^T)   -> opN(Ccm), opT(Vcm)
            //   W2cm = W1cm * T        (=W2^T)   -> opN, opN
            //   Ccm -= W2cm * Vcm      (C -=V W2) -> opN, opN, alpha=-1 beta=1
            float* Cp = Ap + (long long)j * n + jb;       // row-major C base (j,jb)
            long long sC = (long long)n * n, sV = (long long)n * pb,
                      sT = (long long)pb * pb, sW = (long long)pb * n;
            CK_BLAS(cublasSgemmStridedBatched(hbl, CUBLAS_OP_N, CUBLAS_OP_T,
                tcols, pbj, h, &one, Cp, n, sC, Vp, pb, sV, &zero, W1p, tcols, sW, batch));
            CK_BLAS(cublasSgemmStridedBatched(hbl, CUBLAS_OP_N, CUBLAS_OP_N,
                tcols, pbj, pbj, &one, W1p, tcols, sW, Tgp, pb, sT, &zero, W2p, tcols, sW, batch));
            CK_BLAS(cublasSgemmStridedBatched(hbl, CUBLAS_OP_N, CUBLAS_OP_N,
                tcols, h, pbj, &neg, W2p, tcols, sW, Vp, pb, sV, &one, Cp, n, sC, batch));
        }
    }
    CK_CUDA(cudaGetLastError());
    return {At, tau};
}
"""

CPP_SRC = r"""
std::vector<torch::Tensor> qr_compact(torch::Tensor A);
"""

_module = load_inline(
    name="qr_v2_blocked",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["qr_compact"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas", "-lcusolver"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    H, tau = _module.qr_compact(data)
    return (H, tau)