submission 825134

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

_CPP = r"""
#include <pybind11/pybind11.h>
#include <cstdint>
void panel_factor(uintptr_t Hp, uintptr_t taup, int batch, int n, int ps, int pe, int bs, uint64_t smem);
"""

_CU = r"""
#include <cuda_runtime.h>
#include <math.h>
#include <cstdint>

__device__ __forceinline__ float warpAllSum(float v) {
    for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffff, v, o);
    return v;
}

// Factor one panel: columns [ps, pe) over rows [ps, n), in place on H
// (one block per matrix). LAPACK compact-Householder convention (== torch.geqrf):
//   beta = -sign(x0)*||x||  (R diagonal);  tau = (beta - x0)/beta
//   v_i  = A[i,k]/(x0 - beta) for i>k, with implicit v_k = 1
__global__ void panel_kernel(float* __restrict__ H, float* __restrict__ tau, int n, int ps, int pe) {
    int b = blockIdx.x;
    float* A = H + (size_t)b * n * n;
    float* t = tau + (size_t)b * n;
    int tid = threadIdx.x, bs = blockDim.x;

    extern __shared__ float sh[];
    float* v   = sh;        // current reflector (length n)
    float* red = sh + n;    // reduction scratch (length bs)
    __shared__ float s_tau, s_scale;

    for (int k = ps; k < pe; ++k) {
        float part = 0.f;
        for (int i = k + tid; i < n; i += bs) { float val = A[(size_t)i * n + k]; v[i] = val; part += val * val; }
        red[tid] = part; __syncthreads();
        for (int s = bs >> 1; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid + s]; __syncthreads(); }
        float norm_sq = red[0]; __syncthreads();

        float x0 = v[k], norm = sqrtf(norm_sq);
        if (tid == 0) {
            float beta, tau_k, scale;
            if (norm < 1e-30f) { beta = x0; tau_k = 0.f; scale = 0.f; }
            else { beta = (x0 >= 0.f) ? -norm : norm; tau_k = (beta - x0) / beta; scale = 1.f / (x0 - beta); }
            s_tau = tau_k; s_scale = scale; t[k] = tau_k;
            A[(size_t)k * n + k] = beta; v[k] = 1.f;
        }
        __syncthreads();
        float scale = s_scale, tau_k = s_tau;
        for (int i = k + 1 + tid; i < n; i += bs) { float vv = v[i] * scale; v[i] = vv; A[(size_t)i * n + k] = vv; }
        __syncthreads();

        // apply reflector k to the remaining panel columns (k, pe) only.
        // One WARP per column (32 lanes split the rows) so the apply uses the
        // whole block, not just (pe-k-1) threads -- critical for the small-batch
        // large-n panels (n=2048 b=8, n=4096 b=2) that have very few blocks.
        if (tau_k != 0.f) {
            int lane = tid & 31, warp = tid >> 5, nwarps = bs >> 5;
            for (int c = k + 1 + warp; c < pe; c += nwarps) {
                float wsum = 0.f;
                for (int i = k + lane; i < n; i += 32) wsum += v[i] * A[(size_t)i * n + c];
                float wj = warpAllSum(wsum) * tau_k;
                for (int i = k + lane; i < n; i += 32) A[(size_t)i * n + c] -= wj * v[i];
            }
        }
        __syncthreads();
    }
}

void panel_factor(uintptr_t Hp, uintptr_t taup, int batch, int n, int ps, int pe, int bs, uint64_t smem) {
    panel_kernel<<<batch, bs, (size_t)smem>>>((float*)Hp, (float*)taup, n, ps, pe);
}
"""

_MOD = None
def _mod():
    global _MOD
    if _MOD is None:
        _MOD = load_inline(
            name="qr_wy_v5", cpp_sources=_CPP, cuda_sources=_CU, functions=["panel_factor"],
            no_implicit_headers=True, with_pytorch_error_handling=False,
            # sm_90 native (Hopper/H200) + compute_90 PTX so the SAME build JITs
            # forward to the B200 leaderboard (sm_100). A sm_90-only cubin would
            # fail on B200 with "no kernel image is available for execution".
            extra_cuda_cflags=["-O3", "-gencode=arch=compute_90,code=sm_90",
                               "-gencode=arch=compute_90,code=compute_90"],
        )
    return _MOD


# Use true FP32 accumulation in the trailing GEMMs (the correctness gate is fp32
# accuracy over ill-conditioned inputs). TF32 can be enabled later as a routing
# optimization for the well-conditioned majority.
torch.backends.cuda.matmul.allow_tf32 = False


import os


def _panel_width(n: int, batch: int) -> int:
    ov = os.environ.get("QR_NB")
    if ov:
        return int(ov)
    # Small n: one panel (pure on-chip factorization, no GEMM/torch overhead).
    # Larger n: narrow panels so the bulk lands in the GPU-filling trailing GEMM
    # and little work stays on the few-block panel kernel.
    if n <= 64:
        return n
    if n <= 256:
        return 64
    return 16


def _block_threads(n: int, batch: int) -> int:
    ov = os.environ.get("QR_BS")
    if ov:
        return int(ov)
    # Few matrices => more threads per block to cover each SM in the panel kernel.
    return 256 if batch >= 64 else 1024


def custom_kernel(data: input_t) -> output_t:
    A = data.contiguous()
    batch, n, _ = A.shape
    dev = A.device
    H = A.clone()
    tau = torch.zeros((batch, n), device=dev, dtype=torch.float32)
    mod = _mod()
    bs = _block_threads(n, batch)
    nb = _panel_width(n, batch)
    eye_cache = {}

    for ps in range(0, n, nb):
        pe = min(ps + nb, n)
        w = pe - ps
        mod.panel_factor(H.data_ptr(), tau.data_ptr(), batch, n, ps, pe, bs, (n + bs) * 4)
        if pe >= n:
            break
        L = n - ps
        sub = H[:, ps:n, ps:pe]                       # (batch, L, w): R(top) + v(below)
        V = sub.clone()
        top = V[:, :w, :]
        eye_w = eye_cache.get(w)
        if eye_w is None:
            eye_w = torch.eye(w, device=dev, dtype=torch.float32)
            eye_cache[w] = eye_w
        V[:, :w, :] = torch.tril(top, -1) + eye_w     # unit-lower-trapezoidal

        taup = tau[:, ps:pe]                          # (batch, w)
        keep = (taup != 0).to(torch.float32)          # zero out tau=0 reflectors
        Veff = V * keep.unsqueeze(1)
        inv_tau = torch.where(taup != 0, 1.0 / torch.where(taup != 0, taup, torch.ones_like(taup)),
                              torch.ones_like(taup))
        S = Veff.transpose(-1, -2) @ Veff             # (batch, w, w)
        M = torch.triu(S, 1) + torch.diag_embed(inv_tau)
        T = torch.linalg.solve_triangular(M, eye_w.expand(batch, w, w), upper=True, left=True)

        C = H[:, ps:n, pe:n]                          # (batch, L, n-pe)
        Wm = Veff.transpose(-1, -2) @ C               # (batch, w, n-pe)
        Wm = T.transpose(-1, -2) @ Wm
        C -= Veff @ Wm

    return H, tau