submission 801506

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200

"""
Blocked Householder QR (right-looking).

The reference torch.geqrf does the *entire* n-wide factorization inside cusolver
using BLAS-2 rank-1 trailing updates -> no tensor cores, ~112 ms for b640-n512.

Here we block it:
  - factor only a thin b-column PANEL with geqrf (cheap: O(m b^2) vs O(m n^2)),
  - apply the accumulated reflectors to the wide trailing submatrix as ONE
    batched WY update  C -= Y (T^T (Y^T C))  -> three big GEMMs on tensor cores.

geqrf already returns LAPACK compact (H, tau) form, so Y and tau drop straight
out of the panel factorization with no conversion.
"""

import torch
from task import input_t, output_t

BLOCK = 64


def _build_Y(A: torch.Tensor, k: int, b: int) -> torch.Tensor:
    Y = torch.tril(A[:, k:, k:k + b], diagonal=-1).clone()
    idx = torch.arange(b, device=A.device)
    Y[:, idx, idx] = 1.0
    return Y


def _build_T(Y: torch.Tensor, tau: torch.Tensor, k: int, b: int) -> torch.Tensor:
    batch = Y.shape[0]
    T = torch.zeros(batch, b, b, device=Y.device, dtype=Y.dtype)
    G = torch.bmm(Y.mT, Y)  # (batch, b, b) Gram matrix, one GEMM
    for j in range(b):
        T[:, j, j] = tau[:, k + j]
        if j > 0:
            t = -tau[:, k + j: k + j + 1] * G[:, :j, j]
            T[:, :j, j] = torch.bmm(T[:, :j, :j], t.unsqueeze(-1)).squeeze(-1)
    return T


def custom_kernel(data: input_t) -> output_t:
    A = data.clone()
    batch, n, _ = A.shape
    tau = torch.zeros(batch, n, device=A.device, dtype=A.dtype)

    for k in range(0, n, BLOCK):
        b = min(BLOCK, n - k)

        # --- panel factorization (cusolver geqrf on the thin block) ---
        panel = A[:, k:, k:k + b].contiguous()
        Hp, taup = torch.geqrf(panel)
        A[:, k:, k:k + b] = Hp
        tau[:, k:k + b] = taup

        # --- WY trailing update of the wide remainder ---
        if k + b < n:
            Y   = _build_Y(A, k, b)
            T   = _build_T(Y, tau, k, b)
            C   = A[:, k:, k + b:]
            W   = torch.bmm(Y.mT, C)        # Y^T C
            TtW = torch.bmm(T.mT, W)        # T^T Y^T C
            A[:, k:, k + b:] = C - torch.bmm(Y, TtW)

    return A, tau