submission 809249

"""
Batched square compact-Householder QR factorization.

Returns (H, tau) in the same compact convention as torch.geqrf(A):
  - H is batch x n x n FP32: R in the upper triangle, Householder
    vectors v (below the implicit unit) below the diagonal.
  - tau is batch x n FP32: reflector coefficients.

Convention (LAPACK ?geqrf / ?larfg), verified against the checker contract
R_factor - Q.T @ A  with  Q = householder_product(H, tau),  R_factor = triu(H):

  x      = A[k:, k]
  alpha  = x[0]
  xnorm  = || x[1:] ||
  beta   = -sign(alpha) * hypot(alpha, xnorm)
  tau    = (beta - alpha) / beta
  v[0]   = 1,  v[1:] = x[1:] / (alpha - beta)
  H[k,k] = beta   (the diagonal of R)
  H[k+1:, k] = v[1:]
  degenerate xnorm == 0  ->  tau = 0  (column already triangular)

Two execution paths:
  * small n  (<= SMALL_N): a single fused, matrix-resident CUDA kernel,
    one block per matrix, the whole n x n tile kept in shared memory.
  * large n: a blocked compact-WY panel algorithm. A custom CUDA kernel
    factorizes each thin n x nb panel (resident in shared memory) and a
    second kernel forms the nb x nb triangular factor T; the flop-dominant
    trailing-matrix update  A -= V (T^T (V^T A))  is delegated to batched
    cuBLAS GEMM (torch.bmm), which runs at near-peak throughput.

All arithmetic is FP32 end to end; this clears the checker tolerances
(20*n*eps32 factor residual, 100*n*eps32 orthogonality) with wide margin.
"""

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


# --------------------------------------------------------------------------
# CUDA / C++ source
# --------------------------------------------------------------------------
CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#ifndef MAX_THREADS
#define MAX_THREADS 256
#endif

// ---- block reduction helpers (sum) ----------------------------------------
__device__ __forceinline__ float warpReduceSum(float v) {
    for (int o = 16; o > 0; o >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, o);
    return v;
}
__device__ __forceinline__ float blockReduceSum(float v, float* shm) {
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    v = warpReduceSum(v);
    if (lane == 0) shm[wid] = v;
    __syncthreads();
    int nwarp = (blockDim.x + 31) >> 5;
    v = (threadIdx.x < nwarp) ? shm[lane] : 0.0f;
    if (wid == 0) v = warpReduceSum(v);
    // broadcast result via shm[0]
    if (threadIdx.x == 0) shm[0] = v;
    __syncthreads();
    float r = shm[0];
    __syncthreads();
    return r;
}

__device__ __forceinline__ float sgn(float a) {
    // sign with sign(0) = +1, matching LAPACK's choice of -sign for beta
    return (a >= 0.0f) ? 1.0f : -1.0f;
}

// ===========================================================================
// PATH A: fused matrix-resident kernel for small n.
// One block per matrix. The full n x n matrix is loaded column-major into
// shared memory, factorized in place, then written back row-major to H.
// ===========================================================================
extern "C" __global__
void qr_fused_small(const float* __restrict__ A,   // batch x n x n (row-major)
                    float* __restrict__ H,         // batch x n x n (row-major)
                    float* __restrict__ tau,       // batch x n
                    int n) {
    extern __shared__ float smem[];
    float* S    = smem;            // n*n column-major tile: S[i + j*n] = A[i,j]
    float* red  = S + (size_t)n * n;   // reduction scratch (>= warps)

    const int b   = blockIdx.x;
    const int tid = threadIdx.x;
    const int nt  = blockDim.x;
    const float* Ab = A + (size_t)b * n * n;
    float* Hb       = H + (size_t)b * n * n;
    float* taub     = tau + (size_t)b * n;

    // load A (row-major global) into S (column-major shared)
    for (int idx = tid; idx < n * n; idx += nt) {
        int i = idx / n;     // row
        int j = idx % n;     // col
        S[i + j * n] = Ab[idx];
    }
    __syncthreads();

    for (int k = 0; k < n; ++k) {
        float* col = S + (size_t)k * n;   // column k, contiguous in shared

        // norm^2 of x = col[k..n-1]
        float partial = 0.0f;
        for (int i = k + tid; i < n; i += nt) {
            float val = col[i];
            partial += val * val;
        }
        float nrm2 = blockReduceSum(partial, red);
        float alpha = col[k];
        float normx = sqrtf(nrm2);

        // xnorm = ||x[1:]|| ; guard degenerate column
        float tail2 = nrm2 - alpha * alpha;
        if (tail2 < 0.0f) tail2 = 0.0f;

        float beta, tauk, inv;
        if (tail2 <= 0.0f) {
            // already triangular below the diagonal
            beta = alpha;
            tauk = 0.0f;
            inv  = 0.0f;
        } else {
            beta = -sgn(alpha) * normx;
            tauk = (beta - alpha) / beta;
            inv  = 1.0f / (alpha - beta);
        }
        if (tid == 0) taub[k] = tauk;

        // write v in place: v[k] implicit 1 (store beta on diagonal),
        // v[k+1:] = x[k+1:] * inv
        for (int i = k + 1 + tid; i < n; i += nt) col[i] *= inv;
        if (tid == 0) col[k] = beta;
        __syncthreads();

        if (tauk != 0.0f && k + 1 < n) {
            // apply H_k = I - tau v v^T to trailing columns j = k+1..n-1
            // w_j = tau * (v . A[k:,j]) ; A[k:,j] -= v * w_j ; v[k]=1
            for (int j = k + 1 + tid; j < n; j += nt) {
                float* cj = S + (size_t)j * n;
                float dot = cj[k];                 // v[k]=1 contribution
                for (int i = k + 1; i < n; ++i)
                    dot += col[i] * cj[i];
                float w = tauk * dot;
                cj[k] -= w;                         // v[k]=1
                for (int i = k + 1; i < n; ++i)
                    cj[i] -= col[i] * w;
            }
        }
        __syncthreads();
    }

    // write back: S column-major -> H row-major
    for (int idx = tid; idx < n * n; idx += nt) {
        int i = idx / n;
        int j = idx % n;
        Hb[idx] = S[i + j * n];
    }
}

// ===========================================================================
// PATH B, step 1: factorize a thin panel of nb columns for the blocked
// algorithm. Operates on the sub-block H[k:, k:k+nb] of each matrix.
//
// SMEM-bounded: the kernel works directly on the global H matrix (column k+j
// has stride n in row-major storage) and keeps only tiny scratch in shared
// memory, so its shared-memory footprint is independent of n. This is what
// lets it run at n = 4096 where an in-SMEM slab would need megabytes.
//
//   H      : batch x n x n  (row-major, modified in place)
//   Vout   : batch x n x nb (row-major) -- unit-lower Householder vectors,
//            global row index; rows 0..k-1 are zero, row k+j carries the unit
//   tau    : batch x n
//   k      : starting column of the panel
//   nb     : panel width (actual columns this call)
//   n      : matrix dimension
// After this kernel, H[k:, k:k+nb] holds the factored panel (R on/above the
// block diagonal, v below), and Vout holds the explicit V with unit diagonal,
// ready for the WY trailing update via GEMM.
// ===========================================================================
extern "C" __global__
void qr_panel(float* __restrict__ H,
              float* __restrict__ Vout,
              float* __restrict__ tau,
              int k, int nb, int n) {
    extern __shared__ float smem[];
    float* red = smem;                 // reduction scratch (>= warps)

    const int b   = blockIdx.x;
    const int tid = threadIdx.x;
    const int nt  = blockDim.x;
    float* Hb   = H   + (size_t)b * n * n;
    float* Vb   = Vout+ (size_t)b * n * nb;
    float* taub = tau + (size_t)b * n;

    for (int j = 0; j < nb; ++j) {
        const int c   = k + j;             // global column index
        const int top = c;                 // first row of x (the diagonal)
        // pointer arithmetic into column c: element (row r) is Hb[r*n + c]
        // norm^2 of x = H[top..n-1, c]
        float partial = 0.0f;
        for (int r = top + tid; r < n; r += nt) {
            float v = Hb[(size_t)r * n + c];
            partial += v * v;
        }
        float nrm2  = blockReduceSum(partial, red);
        float alpha = Hb[(size_t)top * n + c];
        float normx = sqrtf(nrm2);
        float tail2 = nrm2 - alpha * alpha;
        if (tail2 < 0.0f) tail2 = 0.0f;

        float beta, tauj, inv;
        if (tail2 <= 0.0f) { beta = alpha; tauj = 0.0f; inv = 0.0f; }
        else {
            beta = -sgn(alpha) * normx;
            tauj = (beta - alpha) / beta;
            inv  = 1.0f / (alpha - beta);
        }
        if (tid == 0) taub[c] = tauj;

        // store v below the diagonal in place, beta on the diagonal
        for (int r = top + 1 + tid; r < n; r += nt)
            Hb[(size_t)r * n + c] *= inv;
        if (tid == 0) Hb[(size_t)top * n + c] = beta;
        __syncthreads();

        // within-panel update for columns jj = j+1 .. nb-1 (global col cc=k+jj)
        // v has unit at row 'top'; v[r] = Hb[r*n+c] for r>top.
        for (int jj = j + 1; jj < nb; ++jj) {
            const int cc = k + jj;
            if (tauj == 0.0f) continue;
            // dot = v . H[top..n-1, cc] = H[top,cc] + sum_{r>top} v[r]*H[r,cc]
            float part = (tid == 0) ? Hb[(size_t)top * n + cc] : 0.0f;
            for (int r = top + 1 + tid; r < n; r += nt)
                part += Hb[(size_t)r * n + c] * Hb[(size_t)r * n + cc];
            float dot = blockReduceSum(part, red);
            float w = tauj * dot;
            if (tid == 0) Hb[(size_t)top * n + cc] -= w;   // v[top]=1
            for (int r = top + 1 + tid; r < n; r += nt)
                Hb[(size_t)r * n + cc] -= Hb[(size_t)r * n + c] * w;
            __syncthreads();
        }
    }

    // emit explicit V (unit-lower) into Vb : batch x n x nb
    // zero rows above k
    for (int idx = tid; idx < k * nb; idx += nt) Vb[idx] = 0.0f;
    for (int r = k + tid; r < n; r += nt) {
        int i = r - k;                       // panel row
        for (int j = 0; j < nb; ++j) {
            float vval;
            if (i < j)       vval = 0.0f;
            else if (i == j) vval = 1.0f;
            else             vval = Hb[(size_t)r * n + (k + j)];
            Vb[(size_t)r * nb + j] = vval;
        }
    }
}

// ===========================================================================
// PATH B, step 2: form the nb x nb upper-triangular WY factor T per matrix,
// such that  H_k...H_{k+nb-1} = I - V T V^T.
//   T[:,0] = (tau_0, 0, ...)
//   for j: T[j,j]=tau_j;  T[0:j,j] = -tau_j * T[0:j,0:j] @ (V[:,0:j]^T V[:,j])
// VtV  : batch x nb x nb  =  V^T @ V  (precomputed via bmm)
// tau  : batch x n        (we read tau[k .. k+nb-1])
// Tout : batch x nb x nb  (row-major, upper triangular)
// One block per matrix; nb is small (<=64) so a simple sequential build.
// ===========================================================================
extern "C" __global__
void form_T(const float* __restrict__ VtV,
            const float* __restrict__ tau,
            float* __restrict__ Tout,
            int k, int nb, int n) {
    extern __shared__ float smem[];
    float* T = smem;                 // nb x nb working copy (row-major)
    float* z = T + nb * nb;          // length nb scratch

    const int b   = blockIdx.x;
    const int tid = threadIdx.x;
    const int nt  = blockDim.x;
    const float* Vb = VtV + (size_t)b * nb * nb;
    const float* tb = tau + (size_t)b * n + k;
    float* To       = Tout + (size_t)b * nb * nb;

    for (int idx = tid; idx < nb * nb; idx += nt) T[idx] = 0.0f;
    __syncthreads();

    for (int j = 0; j < nb; ++j) {
        float tauj = tb[j];
        if (tid == 0) T[j * nb + j] = tauj;
        if (j > 0) {
            // z[0:j] = V[:,0:j]^T V[:,j] = VtV[0:j, j]
            for (int i = tid; i < j; i += nt) z[i] = Vb[i * nb + j];
            __syncthreads();
            // T[0:j, j] = -tauj * T[0:j,0:j] @ z[0:j]
            for (int r = tid; r < j; r += nt) {
                float acc = 0.0f;
                for (int c = 0; c < j; ++c)
                    acc += T[r * nb + c] * z[c];
                T[r * nb + j] = -tauj * acc;
            }
        }
        __syncthreads();
    }
    for (int idx = tid; idx < nb * nb; idx += nt) To[idx] = T[idx];
}

// ===========================================================================
// Host launchers. These live in the CUDA (nvcc-compiled) translation unit
// because they use the <<<>>> launch syntax and cudaFuncSetAttribute with a
// kernel symbol -- neither of which the plain C++ host compiler understands.
// ===========================================================================
static int pick_threads(int upper) {
    int t = 32;
    while (t < upper && t < 256) t <<= 1;
    if (t > 256) t = 256;
    return t;
}

void qr_fused_small_launch(torch::Tensor A, torch::Tensor H, torch::Tensor tau, int n) {
    int batch = A.size(0);
    int threads = pick_threads(n);
    size_t shmem = ((size_t)n * n + 64) * sizeof(float);
    cudaFuncSetAttribute((const void*)qr_fused_small,
        cudaFuncAttributeMaxDynamicSharedMemorySize, (int)shmem);
    qr_fused_small<<<batch, threads, shmem>>>(
        A.data_ptr<float>(), H.data_ptr<float>(), tau.data_ptr<float>(), n);
}

void qr_panel_launch(torch::Tensor H, torch::Tensor Vout, torch::Tensor tau,
                     int k, int nb, int n) {
    int batch = H.size(0);
    int m = n - k;
    int threads = pick_threads(m);
    // SMEM is just the reduction scratch -> independent of n.
    size_t shmem = 64 * sizeof(float);
    qr_panel<<<batch, threads, shmem>>>(
        H.data_ptr<float>(), Vout.data_ptr<float>(), tau.data_ptr<float>(),
        k, nb, n);
}

void form_T_launch(torch::Tensor VtV, torch::Tensor tau, torch::Tensor Tout,
                   int k, int nb, int n) {
    int batch = VtV.size(0);
    int threads = pick_threads(nb);
    size_t shmem = ((size_t)nb * nb + nb + 8) * sizeof(float);
    form_T<<<batch, threads, shmem>>>(
        VtV.data_ptr<float>(), tau.data_ptr<float>(), Tout.data_ptr<float>(),
        k, nb, n);
}
"""

# Only plain declarations go to the C++ side; everything that touches CUDA
# launch syntax is compiled by nvcc via cuda_sources above.
CPP_SRC = r"""
#include <torch/extension.h>
void qr_fused_small_launch(torch::Tensor A, torch::Tensor H, torch::Tensor tau, int n);
void qr_panel_launch(torch::Tensor H, torch::Tensor Vout, torch::Tensor tau, int k, int nb, int n);
void form_T_launch(torch::Tensor VtV, torch::Tensor tau, torch::Tensor Tout, int k, int nb, int n);
"""

_module = load_inline(
    name="batched_qr_householder",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["qr_fused_small_launch", "qr_panel_launch", "form_T_launch"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


# --------------------------------------------------------------------------
# Tuning knobs
# --------------------------------------------------------------------------
SMALL_N = 128         # n <= SMALL_N uses the fused matrix-resident path
                      # (n=128 -> 64KB SMEM, safe on every supported GPU)
PANEL_NB = 32          # panel width for the blocked WY path


def _qr_small(A: torch.Tensor) -> output_t:
    batch, n, _ = A.shape
    H = torch.empty_like(A)
    tau = torch.empty((batch, n), dtype=A.dtype, device=A.device)
    Ac = A.contiguous()
    _module.qr_fused_small_launch(Ac, H, tau, n)
    return H, tau


def _qr_blocked(A: torch.Tensor) -> output_t:
    batch, n, _ = A.shape
    H = A.contiguous().clone()
    tau = torch.zeros((batch, n), dtype=A.dtype, device=A.device)

    nb_max = PANEL_NB
    k = 0
    while k < n:
        nb = min(nb_max, n - k)

        # V buffer: batch x n x nb (unit-lower Householder vectors, global rows)
        V = torch.empty((batch, n, nb), dtype=A.dtype, device=A.device)

        # 1) factorize the panel H[:, k:, k:k+nb] in place, emit V and tau
        _module.qr_panel_launch(H, V, tau, k, nb, n)

        rest = n - (k + nb)
        if rest > 0:
            # 2) form T (batch x nb x nb), needs V^T V
            Vk = V[:, k:, :]                       # batch x (n-k) x nb
            VtV = torch.bmm(Vk.transpose(1, 2), Vk)  # batch x nb x nb
            T = torch.empty((batch, nb, nb), dtype=A.dtype, device=A.device)
            _module.form_T_launch(VtV, tau, T, k, nb, n)

            # 3) trailing update:  Trail -= V (T^T (V^T Trail))
            Trail = H[:, k:, k + nb:]              # batch x (n-k) x rest
            W = torch.bmm(Vk.transpose(1, 2), Trail)          # nb x rest
            W = torch.bmm(T.transpose(1, 2), W)               # nb x rest
            Trail.sub_(torch.bmm(Vk, W))          # (n-k) x rest

        k += nb

    return H, tau


def custom_kernel(data: input_t) -> output_t:
    A = data
    if not A.is_cuda:
        A = A.cuda()
    A = A.to(torch.float32).contiguous()
    n = A.shape[-1]
    if A.dim() == 2:
        A = A.unsqueeze(0)
        squeeze = True
    else:
        squeeze = False

    if n <= SMALL_N:
        H, tau = _qr_small(A)
    else:
        H, tau = _qr_blocked(A)

    if squeeze:
        H = H.squeeze(0)
        tau = tau.squeeze(0)
    return H, tau