submission 820377

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200
"""v8-tuned (qr_v2): per-size panel width nb (512->32, 1024->48, 2048->24) minimizes
the within-panel work (~ nb) that dominates each size. v8 (qr_v2): v6 panel + 2048 via blocked panel (28.9ms vs geqrf 76.6); 4096 stays geqrf
(nb too thin for batch=2). v6 (qr_v2): v4 robustness + a faster PADDED WARP-PER-COLUMN panel kernel.
Profiling showed the panel is ~half the 512/1024 cost; v4's panel used <=nb=64 of 256
threads. v6 uses all 8 warps over the panel columns with 32 lanes parallelizing the
row loops via shuffle-reductions, and pads the panel leading dim to nb+1 to eliminate
the shared-memory bank conflict that sank the earlier attempt. FP32 throughout (robust
to rankdef/clustered/nearrank/mixed). 512 appears 4x and 1024 3x in the qr_v2 benchmark,
so this targets 7 of 12 ranked cases. smem 64..224; geqrf for 32 and >=2048.
"""
import torch
from task import input_t, output_t

_CPP = r"""
#include <torch/extension.h>
std::tuple<at::Tensor, at::Tensor> qr_smem(at::Tensor A);
void panel_factor(at::Tensor H, at::Tensor tau, int64_t j, int64_t b);
void build_vt(at::Tensor H, at::Tensor tau, at::Tensor Vout, at::Tensor Tout, int64_t j, int64_t b);
void apply_vw(at::Tensor V, at::Tensor W, at::Tensor H, int64_t j, int64_t b);
void diag2();
"""
_CUDA = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#define NT 256
#define FULL 0xffffffffu

__global__ void qr_smem_kernel(const float* __restrict__ A, float* __restrict__ Hout,
                               float* __restrict__ tau, int n) {
    extern __shared__ float sm[];
    float* As = sm; float* v = sm + (size_t)n * n; float* red = v + n;
    const int mat = blockIdx.x, tid = threadIdx.x;
    const float* Amat = A + (size_t)mat * n * n;
    float* Hmat = Hout + (size_t)mat * n * n; float* taumat = tau + (size_t)mat * n;
    for (int i = tid; i < n * n; i += NT) As[i] = Amat[i];
    __syncthreads();
    for (int k = 0; k < n; ++k) {
        float local = 0.f;
        for (int r = k + 1 + tid; r < n; r += NT) { float x = As[r * n + k]; local += x * x; }
        red[tid] = local; __syncthreads();
        for (int s = NT / 2; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid + s]; __syncthreads(); }
        float sigma = red[0]; __syncthreads();
        float alpha = As[k * n + k], beta, tk, denom;
        if (sigma == 0.f) { beta = alpha; tk = 0.f; denom = 1.f; }
        else { float nf = sqrtf(alpha * alpha + sigma); beta = (alpha >= 0.f) ? -nf : nf; tk = (beta - alpha) / beta; denom = alpha - beta; }
        if (tid == 0) { v[k] = 1.f; As[k * n + k] = beta; taumat[k] = tk; }
        for (int r = k + 1 + tid; r < n; r += NT) v[r] = (sigma == 0.f) ? 0.f : As[r * n + k] / denom;
        __syncthreads();
        for (int r = k + 1 + tid; r < n; r += NT) As[r * n + k] = v[r];
        __syncthreads();
        if (tk != 0.f) for (int jj = k + 1 + tid; jj < n; jj += NT) {
            float w = As[k * n + jj];
            for (int r = k + 1; r < n; ++r) w += v[r] * As[r * n + jj];
            float tw = tk * w; As[k * n + jj] -= tw;
            for (int r = k + 1; r < n; ++r) As[r * n + jj] -= v[r] * tw;
        }
        __syncthreads();
    }
    for (int i = tid; i < n * n; i += NT) Hmat[i] = As[i];
}
std::tuple<at::Tensor, at::Tensor> qr_smem(at::Tensor A) {
    const at::cuda::CUDAGuard guard(A.device());
    A = A.contiguous();
    const int64_t batch = A.size(0); const int n = (int)A.size(1);
    auto H = at::empty_like(A); auto tau = at::empty({batch, n}, A.options());
    size_t shmem = ((size_t)n * n + n + NT) * sizeof(float);
    cudaFuncSetAttribute(qr_smem_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)shmem);
    qr_smem_kernel<<<(int)batch, NT, shmem>>>(A.data_ptr<float>(), H.data_ptr<float>(), tau.data_ptr<float>(), n);
    cudaError_t e = cudaGetLastError(); TORCH_CHECK(e == cudaSuccess, "qr_smem: ", cudaGetErrorString(e));
    return std::make_tuple(H, tau);
}

// Padded warp-per-column panel. ld = b+1 (odd) => smem bank-conflict-free. 8 warps over
// columns; 32 lanes stride rows; shuffle-reduce the within-panel dot products.
__global__ void panel_kernel(float* __restrict__ H, float* __restrict__ tau,
                             int n, int j, int m, int b) {
    extern __shared__ float sm[];
    const int ld = b + 1;
    float* P = sm; float* red = P + (size_t)m * ld;
    const int mat = blockIdx.x, tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
    float* Hmat = H + (size_t)mat * n * n; float* taumat = tau + (size_t)mat * n;
    for (int i = tid; i < m * b; i += NT) { int r = i / b, c = i % b; P[r * ld + c] = Hmat[(size_t)(j + r) * n + (j + c)]; }
    __syncthreads();
    for (int k = 0; k < b; ++k) {
        float local = 0.f;
        for (int r = k + 1 + tid; r < m; r += NT) { float x = P[r * ld + k]; local += x * x; }
        red[tid] = local; __syncthreads();
        for (int s = NT / 2; s > 0; s >>= 1) { if (tid < s) red[tid] += red[tid + s]; __syncthreads(); }
        float sigma = red[0]; __syncthreads();
        float alpha = P[k * ld + k], beta, tk, denom;
        if (sigma == 0.f) { beta = alpha; tk = 0.f; denom = 1.f; }
        else { float nf = sqrtf(alpha * alpha + sigma); beta = (alpha >= 0.f) ? -nf : nf; tk = (beta - alpha) / beta; denom = alpha - beta; }
        if (tid == 0) { P[k * ld + k] = beta; taumat[j + k] = tk; }
        for (int r = k + 1 + tid; r < m; r += NT) P[r * ld + k] = (sigma == 0.f) ? 0.f : P[r * ld + k] / denom;
        __syncthreads();
        if (tk != 0.f) for (int c = k + 1 + warp; c < b; c += 8) {
            float part = 0.f;
            for (int r = k + 1 + lane; r < m; r += 32) part += P[r * ld + k] * P[r * ld + c];
            for (int o = 16; o > 0; o >>= 1) part += __shfl_down_sync(FULL, part, o);
            float w = __shfl_sync(FULL, part, 0) + P[k * ld + c];
            float tw = tk * w;
            if (lane == 0) P[k * ld + c] -= tw;
            for (int r = k + 1 + lane; r < m; r += 32) P[r * ld + c] -= P[r * ld + k] * tw;
        }
        __syncthreads();
    }
    for (int i = tid; i < m * b; i += NT) { int r = i / b, c = i % b; Hmat[(size_t)(j + r) * n + (j + c)] = P[r * ld + c]; }
}
void panel_factor(at::Tensor H, at::Tensor tau, int64_t j, int64_t b) {
    const at::cuda::CUDAGuard guard(H.device());
    int n = (int)H.size(1); int batch = (int)H.size(0); int m = n - (int)j; int ld = (int)b + 1;
    size_t shmem = ((size_t)m * ld + NT) * sizeof(float);
    cudaFuncSetAttribute(panel_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)shmem);
    panel_kernel<<<batch, NT, shmem>>>(H.data_ptr<float>(), tau.data_ptr<float>(), n, (int)j, m, (int)b);
    cudaError_t e = cudaGetLastError();
    TORCH_CHECK(e == cudaSuccess, "panel_kernel: ", cudaGetErrorString(e), " shmem=", (int)shmem);
}

// Fused V+T construction: one CTA per matrix. Replaces the ~10 per-panel torch ops
// (clone/tril/eye/mask/S=VtV/solve_triangular). Reads factored reflectors from H, writes
// clean V (m x b) and the compact-WY T = M^{-1} (b x b, M = striu(VtV) + diag(1/tau)).
__global__ void build_vt_kernel(const float* __restrict__ H, const float* __restrict__ tau,
                                float* __restrict__ Vout, float* __restrict__ Tout,
                                int n, int j, int b, int m) {
    extern __shared__ float sm[];
    float* Vs = sm;                          // m*b  clean V
    float* Ms = Vs + (size_t)m * b;          // b*b  M (upper-tri)
    float* Ts = Ms + (size_t)b * b;          // b*b  T = M^{-1}
    const int mat = blockIdx.x, tid = threadIdx.x;
    const float* Hm = H + (size_t)mat * n * n;
    const float* taum = tau + (size_t)mat * n + j;
    float* Voutm = Vout + (size_t)mat * m * b;
    float* Toutm = Tout + (size_t)mat * b * b;
    // 1. clean V (and copy to Vout)
    for (int i = tid; i < m * b; i += NT) {
        int r = i / b, k = i - r * b;
        float val = (r < k) ? 0.f : (r == k) ? 1.f : Hm[(size_t)(j + r) * n + (j + k)];
        if (taum[k] == 0.f) val = 0.f;
        Vs[i] = val; Voutm[i] = val;
    }
    __syncthreads();
    // 2. M = striu(VtV) + diag(1/tau)
    for (int idx = tid; idx < b * b; idx += NT) {
        int i = idx / b, k = idx - i * b;
        if (i > k) { Ms[idx] = 0.f; }
        else if (i == k) { float t = taum[k]; Ms[idx] = (t != 0.f) ? (1.f / t) : 1.f; }
        else { float s = 0.f; for (int r = 0; r < m; ++r) s += Vs[r * b + i] * Vs[r * b + k]; Ms[idx] = s; }
    }
    __syncthreads();
    // 3. T = M^{-1} (upper-tri), one thread per column jj (validated back-substitution)
    for (int jj = tid; jj < b; jj += NT) {
        Ts[jj * b + jj] = 1.f / Ms[jj * b + jj];
        for (int i = jj - 1; i >= 0; --i) {
            float s = 0.f;
            for (int kk = i + 1; kk <= jj; ++kk) s += Ms[i * b + kk] * Ts[kk * b + jj];
            Ts[i * b + jj] = -s / Ms[i * b + i];
        }
        for (int i = jj + 1; i < b; ++i) Ts[i * b + jj] = 0.f;
    }
    __syncthreads();
    for (int idx = tid; idx < b * b; idx += NT) Toutm[idx] = Ts[idx];
}
void build_vt(at::Tensor H, at::Tensor tau, at::Tensor Vout, at::Tensor Tout, int64_t j, int64_t b) {
    const at::cuda::CUDAGuard guard(H.device());
    int n = (int)H.size(1); int batch = (int)H.size(0); int m = (int)Vout.size(1);
    size_t shmem = ((size_t)m * b + 2 * (size_t)b * b) * sizeof(float);
    cudaFuncSetAttribute(build_vt_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)shmem);
    build_vt_kernel<<<batch, NT, shmem>>>(H.data_ptr<float>(), tau.data_ptr<float>(),
        Vout.data_ptr<float>(), Tout.data_ptr<float>(), n, (int)j, (int)b, m);
    cudaError_t e = cudaGetLastError();
    TORCH_CHECK(e == cudaSuccess, "build_vt: ", cudaGetErrorString(e), " shmem=", (int)shmem);
}

// High-occupancy skinny-K SGEMM (Iter-1 retune): 256 threads, 4x4 micro-tile -> ~75% occ
// (Iter-1's 64-thread/8x8 was 18.75% occ and lost to cuBLAS). C -= V @ W, K=b.
#define VW_BM 64
#define VW_BN 64
#define VW_TM 4
#define VW_TN 4
#define VW_NT 256
#define VW_KMAX 64
__global__ void apply_vw_kernel(const float* __restrict__ V, const float* __restrict__ W,
                                float* __restrict__ H, int n, int j, int b, int m, int ncol) {
    const int mat = blockIdx.z;
    const int row0 = blockIdx.y * VW_BM, col0 = blockIdx.x * VW_BN;
    const float* Vm = V + (size_t)mat * m * b;
    const float* Wm = W + (size_t)mat * b * ncol;
    float* Hm = H + (size_t)mat * n * n;
    __shared__ float As[VW_BM * VW_KMAX];
    __shared__ float Bs[VW_KMAX * VW_BN];
    const int tid = threadIdx.x;
    for (int i = tid; i < VW_BM * b; i += VW_NT) {
        int rr = i / b, kk = i - rr * b, gr = row0 + rr;
        As[rr * b + kk] = (gr < m) ? Vm[(size_t)gr * b + kk] : 0.f;
    }
    for (int i = tid; i < b * VW_BN; i += VW_NT) {
        int kk = i / VW_BN, cc = i - kk * VW_BN, gc = col0 + cc;
        Bs[kk * VW_BN + cc] = (gc < ncol) ? Wm[(size_t)kk * ncol + gc] : 0.f;
    }
    __syncthreads();
    const int trow = (tid / (VW_BN / VW_TN)) * VW_TM;
    const int tcol = (tid % (VW_BN / VW_TN)) * VW_TN;
    float acc[VW_TM][VW_TN];
    #pragma unroll
    for (int a = 0; a < VW_TM; ++a) for (int c = 0; c < VW_TN; ++c) acc[a][c] = 0.f;
    for (int kk = 0; kk < b; ++kk) {
        float ar[VW_TM], br[VW_TN];
        #pragma unroll
        for (int a = 0; a < VW_TM; ++a) ar[a] = As[(trow + a) * b + kk];
        #pragma unroll
        for (int c = 0; c < VW_TN; ++c) br[c] = Bs[kk * VW_BN + tcol + c];
        #pragma unroll
        for (int a = 0; a < VW_TM; ++a) for (int c = 0; c < VW_TN; ++c) acc[a][c] += ar[a] * br[c];
    }
    #pragma unroll
    for (int a = 0; a < VW_TM; ++a) {
        int gr = row0 + trow + a; if (gr >= m) continue;
        #pragma unroll
        for (int c = 0; c < VW_TN; ++c) { int gc = col0 + tcol + c; if (gc < ncol) Hm[(size_t)(j + gr) * n + (j + b + gc)] -= acc[a][c]; }
    }
}
void apply_vw(at::Tensor V, at::Tensor W, at::Tensor H, int64_t j, int64_t b) {
    const at::cuda::CUDAGuard guard(H.device());
    int n = (int)H.size(1), batch = (int)H.size(0), m = (int)V.size(1), ncol = (int)W.size(2);
    if (ncol <= 0) return;
    dim3 grid((ncol + VW_BN - 1) / VW_BN, (m + VW_BM - 1) / VW_BM, batch);
    apply_vw_kernel<<<grid, VW_NT>>>(V.data_ptr<float>(), W.data_ptr<float>(), H.data_ptr<float>(), n, (int)j, (int)b, m, ncol);
    TORCH_CHECK(cudaGetLastError() == cudaSuccess, "apply_vw");
}
void diag2() {
    cudaFuncAttributes a; int occ = 0;
    cudaFuncGetAttributes(&a, (const void*)apply_vw_kernel);
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occ, (const void*)apply_vw_kernel, VW_NT, 0);
    printf("DIAG| apply_vw regs=%d smem=%d activeBlk/SM=%d -> occ=%.0f%%\n",
           a.numRegs, (int)a.sharedSizeBytes, occ, 100.0 * occ * VW_NT / 2048.0);
    fflush(stdout);
}
"""
_MOD = None
try:
    from torch.utils.cpp_extension import load_inline
    _MOD = load_inline(name="qr_p5", cpp_sources=[_CPP], cuda_sources=[_CUDA],
                       functions=["qr_smem", "panel_factor", "build_vt", "apply_vw", "diag2"], extra_cuda_cflags=["-O3"], verbose=False)
except Exception:
    _MOD = None


def _pick_b(n: int) -> int:
    if n == 512:
        return 32          # batch 640: small nb minimizes panel work (trailing stays efficient)
    if n <= 512:
        return 64          # 256/352 (small batch): larger nb keeps trailing GEMM efficient
    if n <= 1024:
        return 48
    if n <= 2048:
        return 24          # smem cap for m=2048 (m*(nb+1)*4 < 227KB)
    b = min(64, (56000 // n))
    return max(8, (b // 8) * 8)


def _qr_panel(A: torch.Tensor) -> output_t:
    H = A.contiguous().clone()
    B, n, _ = H.shape
    tau = torch.zeros(B, n, device=H.device, dtype=H.dtype)
    b0 = _pick_b(n)
    j = 0
    while j < n:
        bb = min(b0, n - j)
        _MOD.panel_factor(H, tau, j, bb)
        if j + bb < n:
            mm = n - j
            V = torch.empty(B, mm, bb, device=H.device, dtype=H.dtype)
            T = torch.empty(B, bb, bb, device=H.device, dtype=H.dtype)
            _MOD.build_vt(H, tau, V, T, j, bb)
            Atrail = H[:, j:, j + bb:]
            W = torch.matmul(V.transpose(1, 2), Atrail)
            W = torch.matmul(T.transpose(1, 2), W)
            _MOD.apply_vw(V, W.contiguous(), H, j, bb)
        j += bb
    return H, tau


def custom_kernel(data: input_t) -> output_t:
    A = data
    n = A.shape[-1]
    if _MOD is not None and A.is_cuda:
        try:
            if 64 <= n <= 224:
                return _MOD.qr_smem(A.contiguous())
            if 256 <= n <= 2048:
                return _qr_panel(A)
        except Exception:
            return torch.geqrf(A)
    return torch.geqrf(A)