submission 843238

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200
"""
qr_v2 — batched square compact-Householder QR.

Attempt 2: hand-written CuTe DSL kernel (no vendor BLAS). Unblocked Householder
QR (geqr2), ONE CTA PER MATRIX, 256 threads/CTA, factored in place in global
memory, matching torch.geqrf's compact (H, tau). The win thesis: cuSOLVER
(torch.geqrf) factors one matrix at a time and underfills the B200; we run all
`batch` matrices concurrently, one CTA each, to fill it.

This pass routes ONLY n==512 (the dominant benchmark shape, batch=640) to the
kernel; every other n falls back to torch.geqrf. SIMT/scalar for now (no tensor
cores) — Phase 2 swaps the trailing update to Blackwell tcgen05 MMA. Any
build/runtime failure falls back to torch.geqrf, so correctness is never at risk.
"""
import torch
from task import input_t, output_t

CUTE_MAXN = 1024  # route n <= this to the CuTe kernel (smem buffer MAXN=1024);
# n in {32,176,352,512,1024} -> kernel; n in {2048,4096} -> geqrf (tiny batch
# can't fill the GPU + exceed the smem reflector buffer).

# NOTE: no runtime pip install — leaderboard-mode KernelGuard blocks
# RUNTIME_PACKAGE_INSTALL. nvidia-cutlass-dsl is preinstalled on the B200 runner.
_CUTE_OK = True
try:
    import cutlass
    import cutlass.cute as cute
    from cutlass.cute.runtime import from_dlpack
    from cutlass.cute import math as cmath

    TPB = 256
    MAXN = 1024  # smem reflector buffer (covers n up to 1024)

    @cute.kernel
    def _geqr2_par_kernel(gH: cute.Tensor, gTau: cute.Tensor):
        tidx, _, _ = cute.arch.thread_idx()
        b, _, _ = cute.arch.block_idx()
        n = gH.shape[1]
        zero = cute.Float32(0.0)
        one = cute.Float32(1.0)

        smem = cutlass.utils.SmemAllocator()
        vs = smem.allocate_array(element_type=cutlass.Float32, num_elems=MAXN)
        red = smem.allocate_array(element_type=cutlass.Float32, num_elems=TPB)

        for k in cutlass.range(n):
            alpha = gH[b, k, k]
            # cooperative column-norm: xnorm2 = sum_{i>k} H[i,k]^2
            partial = zero
            for i in cutlass.range(k + 1 + tidx, n, TPB):
                hik = gH[b, i, k]
                partial = partial + hik * hik
            red[tidx] = partial
            cute.arch.sync_threads()
            if tidx == 0:
                s = zero
                for t in cutlass.range(TPB):
                    s = s + red[t]
                red[0] = s
            cute.arch.sync_threads()
            xnorm2 = red[0]
            cute.arch.sync_threads()

            # reflector scalars (uniform across threads), LAPACK larfg convention
            nrm = cmath.sqrt(alpha * alpha + xnorm2)
            beta = alpha
            d = one
            tau_k = zero
            if xnorm2 > zero:
                ba = nrm
                if alpha >= zero:
                    ba = -nrm
                else:
                    ba = nrm
                beta = ba
                d = alpha - ba
                tau_k = (ba - alpha) / ba

            # form v into smem + store reflector below diagonal
            if tidx == 0:
                vs[k] = one
                gH[b, k, k] = beta
                gTau[b, k] = tau_k
            for i in cutlass.range(k + 1 + tidx, n, TPB):
                vi = gH[b, i, k] / d
                vs[i] = vi
                gH[b, i, k] = vi
            cute.arch.sync_threads()

            # trailing update over columns j>k:  col_j -= tau*v*(v^T col_j)
            for j in cutlass.range(k + 1 + tidx, n, TPB):
                w = gH[b, k, j]
                for i in cutlass.range(k + 1, n):
                    w = w + vs[i] * gH[b, i, j]
                wt = tau_k * w
                gH[b, k, j] = gH[b, k, j] - wt
                for i in cutlass.range(k + 1, n):
                    gH[b, i, j] = gH[b, i, j] - vs[i] * wt
            cute.arch.sync_threads()

    @cute.jit
    def _geqr2_par(mH: cute.Tensor, mTau: cute.Tensor):
        B = mH.shape[0]
        _geqr2_par_kernel(mH, mTau).launch(grid=(B, 1, 1), block=(TPB, 1, 1))

    try:
        cutlass.cuda.initialize_cuda_context()
    except Exception:
        pass

    _CACHE = {}

    def _cute_geqrf(data: torch.Tensor):
        B, n, _ = data.shape
        H = data.clone().contiguous()
        tau = torch.zeros(B, n, device=data.device, dtype=torch.float32)
        H_ = from_dlpack(H)
        tau_ = from_dlpack(tau)
        key = (B, n)
        fn = _CACHE.get(key)
        if fn is None:
            fn = cute.compile(_geqr2_par, H_, tau_)
            _CACHE[key] = fn
        fn(H_, tau_)
        return H, tau

except Exception as exc:  # noqa: BLE001
    print(f"[qr_v2] CuTe DSL unavailable, using torch.geqrf: {exc}", flush=True)
    _CUTE_OK = False


def custom_kernel(data: input_t) -> output_t:
    if _CUTE_OK and data.shape[-1] <= CUTE_MAXN:
        try:
            return _cute_geqrf(data)
        except Exception:  # noqa: BLE001 — never regress correctness
            return torch.geqrf(data)
    return torch.geqrf(data)