submission 831664

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

# This file is GENERATED by submit.sh from qr.cu -- edit qr.cu, not this.

_CUDA_SRC = r"""#include <torch/extension.h>
#include <cusolverDn.h>
#include <cublas_v2.h>
#include <vector>

// Batched compact-Householder QR (LAPACK geqrf convention), matching torch.geqrf:
//   - upper triangle of H holds R, below-diagonal holds the reflector vectors v
//   - tau holds the reflector scalar factors
//
// torch.geqrf loops cuSOLVER one matrix at a time; a single small/mid solve
// underutilizes the GPU, so the serial batch wastes most of the machine. For
// mid n we instead run a blocked (compact-WY) batched QR on one queue:
//   - a custom batched kernel factors a narrow panel of NB columns for every
//     matrix at once (cheap BLAS-2, one block per matrix) and builds, per
//     matrix, the panel reflectors V and the NB x NB WY factor T;
//   - the trailing-matrix update A_trail -= V (T^T (V^T A_trail)) is three
//     cuBLAS strided-batched GEMMs (BLAS-3) that fill the GPU across the batch.
// The block reflector is mathematically the sequential product of the panel's
// Householder reflectors, so R and tau match LAPACK exactly.
//
// Other shapes fall back to looped cuSOLVER (a lone large matrix already fills
// the GPU; a lone small matrix has too little work to amortize the panels).

#define CUSOLVER_CHECK(expr) \
    TORCH_CHECK((expr) == CUSOLVER_STATUS_SUCCESS, "cuSOLVER error at " #expr)
#define CUBLAS_CHECK(expr) \
    TORCH_CHECK((expr) == CUBLAS_STATUS_SUCCESS, "cuBLAS error at " #expr)

#define NB 8             // panel width
#define NTHREADS 256
#define BLOCKED_MIN_N 384
#define BLOCKED_MAX_N 1536

static cusolverDnHandle_t cusolver_handle() {
    static cusolverDnHandle_t h = nullptr;
    if (!h) CUSOLVER_CHECK(cusolverDnCreate(&h));
    return h;
}
static cublasHandle_t cublas_handle() {
    static cublasHandle_t h = nullptr;
    if (!h) CUBLAS_CHECK(cublasCreate(&h));
    return h;
}

// One block per matrix. Column-major A (lda = m): A(i,j) = Acol[(long)j*m + i].
// Factors columns [c0, c0+pb), writes the masked unit-lower-trapezoidal panel
// into Vbuf (col-major m x NB, lda = m) and the WY factor into Tbuf (col-major
// NB x NB, lda = NB).
__global__ void panel_kernel(float* __restrict__ Acol, float* __restrict__ tau,
                             float* __restrict__ Vbuf, float* __restrict__ Tbuf,
                             int m, int n, int c0, int pb) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int nt = blockDim.x;
    float* Ab = Acol + (long)b * m * n;
    int k = (m < n) ? m : n;
    float* taub = tau + (long)b * k;
    float* Vb = Vbuf + (long)b * NB * m;
    float* Tb = Tbuf + (long)b * NB * NB;
    int h = m - c0;

    extern __shared__ char sraw[];
    float* red = reinterpret_cast<float*>(sraw);  // [nt]
    float* P = red + nt;                          // [pb*h] panel cache, P[jj*h + lr]
    __shared__ float s_tau, s_beta, s_alpha;

    // Load the active panel (rows [c0,m), cols [c0,c0+pb)) into shared. Local row
    // lr = i - c0; the pivot of panel column jj sits at lr == jj.
    for (int idx = tid; idx < pb * h; idx += nt) {
        int jj = idx / h, lr = idx % h;
        P[(long)jj * h + lr] = Ab[(long)(c0 + jj) * m + (c0 + lr)];
    }
    __syncthreads();

    // ---- factor the pb panel columns in shared (right-looking, within panel) ----
    for (int jj = 0; jj < pb; ++jj) {
        float* Pj = P + (long)jj * h;  // column jj: Pj[lr] = A(c0+lr, c0+jj)

        float local = 0;
        for (int lr = jj + 1 + tid; lr < h; lr += nt) { float v = Pj[lr]; local += v * v; }
        red[tid] = local;
        __syncthreads();
        for (int s = nt / 2; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid + s]; __syncthreads(); }
        if (tid == 0) {
            float sumsq = red[0];
            float xnorm = sqrtf(sumsq);
            float alpha = Pj[jj];
            float tj, beta;
            if (xnorm == 0.f) { tj = 0.f; beta = alpha; }
            else { float r = sqrtf(alpha * alpha + xnorm * xnorm); beta = (alpha >= 0.f) ? -r : r; tj = (beta - alpha) / beta; }
            s_tau = tj; s_beta = beta; s_alpha = alpha; taub[c0 + jj] = tj;
        }
        __syncthreads();
        float tj = s_tau, beta = s_beta, alpha = s_alpha;

        if (tj != 0.f) {
            float scale = 1.f / (alpha - beta);
            for (int lr = jj + 1 + tid; lr < h; lr += nt) Pj[lr] *= scale;
        }
        __syncthreads();
        if (tid == 0) Pj[jj] = beta;
        __syncthreads();

        if (tj != 0.f) {
            for (int col = jj + 1; col < pb; ++col) {
                float* Pc = P + (long)col * h;
                float part = 0;
                for (int lr = jj + 1 + tid; lr < h; lr += nt) part += Pj[lr] * Pc[lr];
                red[tid] = part;
                __syncthreads();
                for (int s = nt / 2; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid + s]; __syncthreads(); }
                float w = (red[0] + Pc[jj]) * tj;
                __syncthreads();
                if (tid == 0) Pc[jj] -= w;
                for (int lr = jj + 1 + tid; lr < h; lr += nt) Pc[lr] -= w * Pj[lr];
                __syncthreads();
            }
        }
    }

    // Write the factored panel (R + reflectors) back to global A.
    for (int idx = tid; idx < pb * h; idx += nt) {
        int jj = idx / h, lr = idx % h;
        Ab[(long)(c0 + jj) * m + (c0 + lr)] = P[(long)jj * h + lr];
    }
    __syncthreads();

    // ---- masked V: Vb(r,jj) = 1 (r==jj), v (r>jj), 0 (r<jj) ----
    for (int idx = tid; idx < h * pb; idx += nt) {
        int r = idx % h, jj = idx / h;
        float val;
        if (r == jj) val = 1.f;
        else if (r > jj) val = P[(long)jj * h + r];
        else val = 0.f;
        Vb[(long)jj * m + r] = val;
    }
    __syncthreads();

    // ---- WY factor T (pb x pb upper triangular), col-major Tb[c*NB + r]=T(r,c) ----
    for (int idx = tid; idx < pb * pb; idx += nt) { int c = idx / pb, r = idx % pb; Tb[(long)c * NB + r] = 0.f; }
    __syncthreads();
    for (int i = 0; i < pb; ++i) {
        float taui = taub[c0 + i];
        // u(r) = -tau_i * sum_{g=i}^{h-1} Vb(g,r) Vb(g,i),  r in [0,i)
        for (int r = tid; r < i; r += nt) {
            float s = 0;
            for (int g = i; g < h; ++g) s += Vb[(long)r * m + g] * Vb[(long)i * m + g];
            Tb[(long)i * NB + r] = -taui * s;
        }
        __syncthreads();
        // T(0:i,i) = T(0:i,0:i) * u   (upper-triangular T times u)
        if (tid < i) red[tid] = Tb[(long)i * NB + tid];   // copy u out of T column i
        __syncthreads();
        for (int r = tid; r < i; r += nt) {
            float s = 0;
            for (int sidx = r; sidx < i; ++sidx) s += Tb[(long)sidx * NB + r] * red[sidx];
            Tb[(long)i * NB + r] = s;
        }
        __syncthreads();
        if (tid == 0) Tb[(long)i * NB + i] = taui;
        __syncthreads();
    }
}

static std::vector<torch::Tensor> geqrf_blocked(torch::Tensor in, int m, int n, long batch) {
    int k = m < n ? m : n;
    auto in3 = in.reshape({batch, m, n});
    auto Acol = in3.transpose(1, 2).contiguous();  // col-major m x n, lda = m
    auto tau = torch::empty({batch, k}, Acol.options());
    auto opts = Acol.options();
    auto Vbuf = torch::empty({batch, NB, m}, opts);   // col-major m x NB, lda = m
    auto Tbuf = torch::empty({batch, NB, NB}, opts);  // col-major NB x NB, lda = NB
    auto Wbuf = torch::empty({batch, n, NB}, opts);   // col-major NB x n, lda = NB
    auto W2buf = torch::empty({batch, n, NB}, opts);

    float* A = Acol.data_ptr<float>();
    float* V = Vbuf.data_ptr<float>();
    float* T = Tbuf.data_ptr<float>();
    float* W = Wbuf.data_ptr<float>();
    float* W2 = W2buf.data_ptr<float>();
    cublasHandle_t handle = cublas_handle();
    const float one = 1.f, zero = 0.f, negone = -1.f;

    // Largest panel cache is at c0 = 0 (pb = NB, h = m). Opt into the shared
    // memory it needs (one block per SM here, so the larger footprint is free).
    size_t maxshmem = ((size_t)NTHREADS + (size_t)NB * m) * sizeof(float);
    cudaFuncSetAttribute((const void*)panel_kernel,
        cudaFuncAttributeMaxDynamicSharedMemorySize, (int)maxshmem);

    for (int c0 = 0; c0 < k; c0 += NB) {
        int pb = (NB < k - c0) ? NB : (k - c0);
        int h = m - c0;
        size_t shmem = ((size_t)NTHREADS + (size_t)pb * h) * sizeof(float);
        panel_kernel<<<batch, NTHREADS, shmem>>>(A, tau.data_ptr<float>(), V, T, m, n, c0, pb);

        int tc = n - (c0 + pb);
        if (tc > 0) {
            float* C = A + (long)(c0 + pb) * m + c0;
            // W = V^T C   (pb x tc)
            CUBLAS_CHECK(cublasSgemmStridedBatched(
                handle, CUBLAS_OP_T, CUBLAS_OP_N, pb, tc, h,
                &one, V, m, (long long)NB * m, C, m, (long long)m * n,
                &zero, W, NB, (long long)n * NB, (int)batch));
            // W2 = T^T W   (pb x tc)
            CUBLAS_CHECK(cublasSgemmStridedBatched(
                handle, CUBLAS_OP_T, CUBLAS_OP_N, pb, tc, pb,
                &one, T, NB, (long long)NB * NB, W, NB, (long long)n * NB,
                &zero, W2, NB, (long long)n * NB, (int)batch));
            // C -= V W2   (h x tc)
            CUBLAS_CHECK(cublasSgemmStridedBatched(
                handle, CUBLAS_OP_N, CUBLAS_OP_N, h, tc, pb,
                &negone, V, m, (long long)NB * m, W2, NB, (long long)n * NB,
                &one, C, m, (long long)m * n, (int)batch));
        }
    }

    auto H = Acol.transpose(1, 2).contiguous();
    return {H, tau};
}

static std::vector<torch::Tensor> geqrf_cusolver(torch::Tensor in, int m, int n, long batch) {
    int k = m < n ? m : n;
    auto in3 = in.reshape({batch, m, n});
    auto Acol = in3.transpose(1, 2).contiguous();
    auto tau = torch::empty({batch, k}, Acol.options());

    cusolverDnHandle_t handle = cusolver_handle();
    int lwork = 0;
    CUSOLVER_CHECK(cusolverDnSgeqrf_bufferSize(handle, m, n, Acol.data_ptr<float>(), m, &lwork));
    auto work = torch::empty({lwork}, Acol.options());
    auto info = torch::empty({1}, torch::dtype(torch::kInt32).device(Acol.device()));

    long Astride = (long)m * n;
    for (long b = 0; b < batch; ++b) {
        CUSOLVER_CHECK(cusolverDnSgeqrf(
            handle, m, n, Acol.data_ptr<float>() + b * Astride, m,
            tau.data_ptr<float>() + b * k, work.data_ptr<float>(), lwork, info.data_ptr<int>()));
    }
    auto H = Acol.transpose(1, 2).contiguous();
    return {H, tau};
}

std::vector<torch::Tensor> geqrf_cuda(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.dim() >= 2, "input must be at least 2D");
    TORCH_CHECK(input.scalar_type() == torch::kFloat32, "input must be float32");

    auto in = input.contiguous();
    int n = in.size(-1);
    int m = in.size(-2);
    TORCH_CHECK(m == n, "only square matrices are supported");
    long batch = in.numel() / ((long)m * n);
    int k = m < n ? m : n;

    std::vector<torch::Tensor> r;
    if (n >= BLOCKED_MIN_N && n <= BLOCKED_MAX_N && batch >= 2) {
        r = geqrf_blocked(in, m, n, batch);
    } else {
        r = geqrf_cusolver(in, m, n, batch);
    }

    auto H = r[0].reshape(input.sizes());
    auto tau_shape = std::vector<int64_t>(input.sizes().begin(), input.sizes().end() - 2);
    tau_shape.push_back(k);
    auto tau_out = r[1].reshape(tau_shape);
    return {H, tau_out};
}
"""

_module = torch.utils.cpp_extension.load_inline(
    name="qr_geqrf",
    cpp_sources="std::vector<at::Tensor> geqrf_cuda(at::Tensor input);",
    cuda_sources=_CUDA_SRC,
    functions=["geqrf_cuda"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcusolver", "-lcublas"],
    verbose=False,
)


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