submission 830792

"""Hybrid QR v4: shared-memory panel kernel (coalesced) + fp32 trailing.

v3 profiling: warp-parallel panel is memory-bound for n>=512 (panel ~76% of time)
because column access A[:,c] is strided (uncoalesced) in row-major. v4 loads the
whole panel A[j:n, j:j+b] into shared memory with COALESCED row-slice loads
(A[r, j:j+b] is contiguous), factors entirely in smem (the arithmetic is tiny;
the cost was memory stalls), and writes back coalesced. nb chosen so the panel
fits in B200's 228KB smem: 64 for n<=512, 32 for n=1024. n<=32 / n>=2048 -> geqrf.
"""
import torch

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

_EXT = None


def _nb_for(n):
    return 32


def _load_ext():
    global _EXT
    if _EXT is not None:
        return _EXT
    from torch.utils.cpp_extension import load_inline
    cuda_src = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#define FULL 0xffffffffu

__global__ void panel_kernel(float* __restrict__ H, float* __restrict__ tau,
                             int B, int n, int j, int b, int ntau, int L) {
    int mat = blockIdx.x;
    float* A = H + (long long)mat * n * n;
    float* tv = tau + (long long)mat * ntau;
    extern __shared__ float sm[];
    float* sp  = sm;          // [L*b] panel (row-major: sp[r*b+c])
    float* wsc = sm + L * b;  // [nwarps]
    int tid = threadIdx.x, nt = blockDim.x;
    int lane = tid & 31, wid = tid >> 5, nwarps = nt >> 5;

    // coalesced load of panel into smem
    for (int e = tid; e < L * b; e += nt) {
        int r = e / b, c = e - r * b;
        sp[e] = A[(long long)(j + r) * n + (j + c)];
    }
    __syncthreads();

    for (int k = 0; k < b; ++k) {
        float alpha = sp[k * b + k];
        float part = 0.f;
        for (int r = k + 1 + tid; r < L; r += nt) { float v = sp[r * b + k]; part += v * v; }
        for (int o = 16; o > 0; o >>= 1) part += __shfl_down_sync(FULL, part, o);
        if (lane == 0) wsc[wid] = part;
        __syncthreads();
        float sumsq = 0.f;
        for (int i = 0; i < nwarps; ++i) sumsq += wsc[i];
        float normx = sqrtf(alpha * alpha + sumsq);
        int safe = normx > 0.f;
        float beta, tau_k, denom;
        if (safe) {
            float sign = (alpha >= 0.f) ? 1.f : -1.f;
            beta = -sign * normx; tau_k = (beta - alpha) / beta; denom = alpha - beta;
        } else { beta = alpha; tau_k = 0.f; denom = 1.f; }
        for (int r = k + 1 + tid; r < L; r += nt) {
            float val = sp[r * b + k];
            sp[r * b + k] = safe ? (val / denom) : val;
        }
        __syncthreads();
        if (tid == 0) { sp[k * b + k] = beta; tv[j + k] = tau_k; }
        __syncthreads();
        // apply reflector to remaining panel cols; warp wid owns column c
        for (int c = k + 1 + wid; c < b; c += nwarps) {
            float wp = 0.f;
            for (int r = k + lane; r < L; r += 32) {
                float vr = (r == k) ? 1.f : sp[r * b + k];
                wp += vr * sp[r * b + c];
            }
            for (int o = 16; o > 0; o >>= 1) wp += __shfl_xor_sync(FULL, wp, o);
            float w = wp * tau_k;
            for (int r = k + lane; r < L; r += 32) {
                float vr = (r == k) ? 1.f : sp[r * b + k];
                sp[r * b + c] -= vr * w;
            }
        }
        __syncthreads();
    }
    // write back (coalesced)
    for (int e = tid; e < L * b; e += nt) {
        int r = e / b, c = e - r * b;
        A[(long long)(j + r) * n + (j + c)] = sp[e];
    }
}

void panel_factor(torch::Tensor H, torch::Tensor tau, int j, int b) {
    int B = H.size(0); int n = H.size(1); int ntau = tau.size(1);
    int L = n - j;
    int threads = 512;
    size_t smem = (size_t)(L * b + threads / 32) * sizeof(float);
    cudaFuncSetAttribute(panel_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
    panel_kernel<<<B, threads, smem>>>(
        H.data_ptr<float>(), tau.data_ptr<float>(), B, n, j, b, ntau, L);
}
'''
    cpp_src = "void panel_factor(torch::Tensor H, torch::Tensor tau, int j, int b);"
    _EXT = load_inline(
        name="qr_panel_v5",
        cpp_sources=cpp_src,
        cuda_sources=cuda_src,
        functions=["panel_factor"],
        extra_cuda_cflags=["-O3"],
        verbose=False,
    )
    return _EXT


def _form_T(V, tau):
    b = V.shape[2]
    eye = torch.eye(b, dtype=V.dtype, device=V.device)
    G = torch.bmm(V.transpose(1, 2), V)
    U = torch.triu(G, diagonal=1)
    D = tau[:, :, None] * eye
    M = eye + torch.bmm(D, U)
    return torch.linalg.solve(M, D)


def _hybrid(A):
    B, n, _ = A.shape
    nb = _nb_for(n)
    ext = _load_ext()
    H = A.contiguous().clone()
    tau = torch.zeros(B, n, dtype=A.dtype, device=A.device)
    eye = torch.eye(nb, dtype=A.dtype, device=A.device)
    j = 0
    while j < n:
        b = min(nb, n - j)
        ext.panel_factor(H, tau, j, b)
        if j + b < n:
            panel = H[:, j:, j:j + b]
            V = torch.tril(panel, diagonal=-1)
            V[:, :b, :b] = V[:, :b, :b] + eye[:b, :b]
            T = _form_T(V, tau[:, j:j + b])
            C = H[:, j:, j + b:]
            W = torch.bmm(V.transpose(1, 2), C)
            W = torch.bmm(T.transpose(1, 2), W)
            C.sub_(torch.bmm(V, W))
        j += b
    return H, tau


def _blocked_cpu(A, nb=64):
    B, n, _ = A.shape
    H = A.clone()
    tau = torch.zeros(B, n, dtype=A.dtype, device=A.device)
    one = torch.ones(B, 1, dtype=A.dtype, device=A.device)
    eye = torch.eye(nb, dtype=A.dtype, device=A.device)
    j = 0
    while j < n:
        b = min(nb, n - j)
        for k in range(j, j + b):
            x = H[:, k:, k]
            alpha = x[:, 0]
            xnorm = torch.linalg.vector_norm(x[:, 1:], dim=1) if x.shape[1] > 1 else torch.zeros_like(alpha)
            normx = torch.sqrt(alpha * alpha + xnorm * xnorm)
            sign = torch.where(alpha >= 0, 1.0, -1.0).to(H.dtype)
            beta = -sign * normx
            safe = normx > 0
            betas = torch.where(safe, beta, torch.ones_like(beta))
            tau_k = torch.where(safe, (beta - alpha) / betas, torch.zeros_like(beta))
            denom = torch.where(safe, alpha - beta, torch.ones_like(beta))
            if x.shape[1] > 1:
                vt = x[:, 1:] / denom[:, None]
                H[:, k + 1:, k] = torch.where(safe[:, None], vt, torch.zeros_like(vt))
            H[:, k, k] = torch.where(safe, beta, alpha)
            tau[:, k] = tau_k
            if k + 1 < j + b:
                sub = H[:, k:, k + 1:j + b]
                v = torch.cat([one[:, :1], H[:, k + 1:, k]], dim=1)
                w = torch.einsum('bi,bij->bj', v, sub)
                sub.sub_((tau_k[:, None] * v)[:, :, None] * w[:, None, :])
        if j + b < n:
            panel = H[:, j:, j:j + b]
            V = torch.tril(panel, diagonal=-1)
            V[:, :b, :b] = V[:, :b, :b] + eye[:b, :b]
            T = _form_T(V, tau[:, j:j + b])
            C = H[:, j:, j + b:]
            W = torch.einsum('bik,bij->bkj', V, C)
            W = torch.bmm(T.transpose(1, 2), W)
            C.sub_(torch.einsum('bik,bkj->bij', V, W))
        j += b
    return H, tau


def custom_kernel(data):
    A = data
    B, n, _ = A.shape
    if not A.is_cuda:
        return _blocked_cpu(A)
    if n <= 32 or n >= 2048:
        return torch.geqrf(A)
    try:
        return _hybrid(A)
    except Exception:
        return torch.geqrf(A)