submission 797901

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200

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


CPP_SRC = """
#include <torch/extension.h>

std::vector<torch::Tensor> qr_householder_cuda(torch::Tensor input);
"""


_qr_module = load_inline(
    name="qr_householder_b200_v5",
    cpp_sources=[CPP_SRC],
    cuda_sources=[
        r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <climits>
#include <math.h>
#include <vector>

namespace {

constexpr int QR_THREADS = 256;

__global__ void row_to_col_kernel(
    const float* __restrict__ src,
    float* __restrict__ dst,
    long long total,
    int n
) {
    long long linear = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (linear >= total) {
        return;
    }

    long long matrix_elems = static_cast<long long>(n) * n;
    long long in_matrix = linear % matrix_elems;
    int row = static_cast<int>(in_matrix / n);
    int col = static_cast<int>(in_matrix - static_cast<long long>(row) * n);
    long long batch = linear / matrix_elems;

    dst[batch * matrix_elems + static_cast<long long>(col) * n + row] = src[linear];
}

__global__ void factor_column_kernel(float* __restrict__ a, float* __restrict__ tau, int n, int k) {
    extern __shared__ float shared[];
    int b = blockIdx.x;
    int tid = threadIdx.x;
    long long matrix_elems = static_cast<long long>(n) * n;
    float* mat = a + static_cast<long long>(b) * matrix_elems;

    float sigma = 0.0f;
    for (int row = k + 1 + tid; row < n; row += blockDim.x) {
        float value = mat[static_cast<long long>(k) * n + row];
        sigma += value * value;
    }

    shared[tid] = sigma;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            shared[tid] += shared[tid + stride];
        }
        __syncthreads();
    }

    __shared__ float tau_k;
    __shared__ float denom;

    if (tid == 0) {
        float alpha = mat[static_cast<long long>(k) * n + k];
        if (shared[0] == 0.0f) {
            tau_k = 0.0f;
            denom = 1.0f;
        } else {
            float norm = sqrtf(alpha * alpha + shared[0]);
            float beta = (alpha >= 0.0f) ? -norm : norm;
            tau_k = (beta - alpha) / beta;
            denom = alpha - beta;
            mat[static_cast<long long>(k) * n + k] = beta;
        }
        tau[static_cast<long long>(b) * n + k] = tau_k;
    }
    __syncthreads();

    if (tau_k == 0.0f) {
        for (int row = k + 1 + tid; row < n; row += blockDim.x) {
            mat[static_cast<long long>(k) * n + row] = 0.0f;
        }
    } else {
        for (int row = k + 1 + tid; row < n; row += blockDim.x) {
            mat[static_cast<long long>(k) * n + row] /= denom;
        }
    }
}

__global__ void apply_householder_kernel(
    float* __restrict__ a,
    const float* __restrict__ tau,
    int n,
    int k,
    int columns
) {
    extern __shared__ float shared[];
    int block = blockIdx.x;
    int b = block / columns;
    int col = k + 1 + (block - b * columns);
    int tid = threadIdx.x;

    float tau_k = tau[static_cast<long long>(b) * n + k];
    if (tau_k == 0.0f) {
        return;
    }

    long long matrix_elems = static_cast<long long>(n) * n;
    float* mat = a + static_cast<long long>(b) * matrix_elems;
    float* v_col = mat + static_cast<long long>(k) * n;
    float* target_col = mat + static_cast<long long>(col) * n;

    float dot = 0.0f;
    for (int row = k + tid; row < n; row += blockDim.x) {
        float v = (row == k) ? 1.0f : v_col[row];
        dot += v * target_col[row];
    }

    shared[tid] = dot;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            shared[tid] += shared[tid + stride];
        }
        __syncthreads();
    }

    float scale = tau_k * shared[0];
    for (int row = k + tid; row < n; row += blockDim.x) {
        float v = (row == k) ? 1.0f : v_col[row];
        target_col[row] -= scale * v;
    }
}

__global__ void apply_householder4_kernel(
    float* __restrict__ a,
    const float* __restrict__ tau,
    int n,
    int k,
    int columns,
    int column_groups
) {
    extern __shared__ float shared[];
    int block = blockIdx.x;
    int b = block / column_groups;
    int group = block - b * column_groups;
    int col0 = k + 1 + group * 4;
    int valid = columns - group * 4;
    valid = valid > 4 ? 4 : valid;
    int tid = threadIdx.x;

    float tau_k = tau[static_cast<long long>(b) * n + k];
    if (tau_k == 0.0f) {
        return;
    }

    long long matrix_elems = static_cast<long long>(n) * n;
    float* mat = a + static_cast<long long>(b) * matrix_elems;
    float* v_col = mat + static_cast<long long>(k) * n;
    float* target0 = mat + static_cast<long long>(col0) * n;
    float* target1 = target0 + static_cast<long long>(n);
    float* target2 = target1 + static_cast<long long>(n);
    float* target3 = target2 + static_cast<long long>(n);

    float dot0 = 0.0f;
    float dot1 = 0.0f;
    float dot2 = 0.0f;
    float dot3 = 0.0f;

    for (int row = k + tid; row < n; row += blockDim.x) {
        float v = (row == k) ? 1.0f : v_col[row];
        dot0 += v * target0[row];
        if (valid > 1) {
            dot1 += v * target1[row];
        }
        if (valid > 2) {
            dot2 += v * target2[row];
        }
        if (valid > 3) {
            dot3 += v * target3[row];
        }
    }

    int base = tid * 4;
    shared[base + 0] = dot0;
    shared[base + 1] = dot1;
    shared[base + 2] = dot2;
    shared[base + 3] = dot3;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) {
            int dst = tid * 4;
            int src = (tid + stride) * 4;
            shared[dst + 0] += shared[src + 0];
            shared[dst + 1] += shared[src + 1];
            shared[dst + 2] += shared[src + 2];
            shared[dst + 3] += shared[src + 3];
        }
        __syncthreads();
    }

    float scale0 = tau_k * shared[0];
    float scale1 = tau_k * shared[1];
    float scale2 = tau_k * shared[2];
    float scale3 = tau_k * shared[3];

    for (int row = k + tid; row < n; row += blockDim.x) {
        float v = (row == k) ? 1.0f : v_col[row];
        target0[row] -= scale0 * v;
        if (valid > 1) {
            target1[row] -= scale1 * v;
        }
        if (valid > 2) {
            target2[row] -= scale2 * v;
        }
        if (valid > 3) {
            target3[row] -= scale3 * v;
        }
    }
}

void launch_row_to_col(const torch::Tensor& src, torch::Tensor& dst, int batch, int n) {
    long long total = static_cast<long long>(batch) * n * n;
    int blocks = static_cast<int>((total + QR_THREADS - 1) / QR_THREADS);
    row_to_col_kernel<<<blocks, QR_THREADS>>>(src.data_ptr<float>(), dst.data_ptr<float>(), total, n);
}

}  // namespace

std::vector<torch::Tensor> qr_householder_cuda(torch::Tensor input) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be torch.float32");
    TORCH_CHECK(input.dim() == 3, "input must have shape batch x n x n");
    TORCH_CHECK(input.size(1) == input.size(2), "input matrices must be square");

    const auto batch64 = input.size(0);
    const auto n64 = input.size(1);
    TORCH_CHECK(batch64 > 0, "batch must be positive");
    TORCH_CHECK(n64 > 0, "n must be positive");
    TORCH_CHECK(batch64 <= static_cast<int64_t>(INT_MAX), "batch is too large");
    TORCH_CHECK(n64 <= static_cast<int64_t>(INT_MAX), "n is too large");

    int batch = static_cast<int>(batch64);
    int n = static_cast<int>(n64);

    c10::cuda::CUDAGuard device_guard(input.device());

    torch::Tensor input_contig = input.contiguous();
    torch::Tensor h = torch::empty_strided(
        {batch64, n64, n64},
        {n64 * n64, 1, n64},
        input_contig.options()
    );
    torch::Tensor tau = torch::empty({batch64, n64}, input_contig.options());

    launch_row_to_col(input_contig, h, batch, n);

    size_t shared_bytes = static_cast<size_t>(QR_THREADS) * sizeof(float);
    for (int k = 0; k < n; ++k) {
        factor_column_kernel<<<batch, QR_THREADS, shared_bytes>>>(
            h.data_ptr<float>(),
            tau.data_ptr<float>(),
            n,
            k
        );
        int columns = n - k - 1;
        if (columns > 0) {
            if (n <= 1024) {
                int column_groups = (columns + 3) / 4;
                apply_householder4_kernel<<<batch * column_groups, QR_THREADS, shared_bytes * 4>>>(
                    h.data_ptr<float>(),
                    tau.data_ptr<float>(),
                    n,
                    k,
                    columns,
                    column_groups
                );
            } else {
                apply_householder_kernel<<<batch * columns, QR_THREADS, shared_bytes>>>(
                    h.data_ptr<float>(),
                    tau.data_ptr<float>(),
                    n,
                    k,
                    columns
                );
            }
        }
    }

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "qr_householder_cuda failed: ", cudaGetErrorString(err));
    return {h, tau};
}
"""
    ],
    functions=["qr_householder_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


_FACTOR_RTOL_FACTOR = 20.0
_ORTH_RTOL_FACTOR = 100.0


def _apply_column_scaling(a: torch.Tensor, cond: int) -> torch.Tensor:
    # `cond` is a deterministic dynamic-range knob, not an exact condition number.
    if cond:
        n = a.shape[-1]
        scales = torch.logspace(0.0, -float(cond), n, device=a.device, dtype=torch.float32)
        return a * scales
    return a.contiguous()


def _band_mask(n: int, bandwidth: int, device: torch.device) -> torch.Tensor:
    idx = torch.arange(n, device=device)
    return (idx[:, None] - idx[None, :]).abs() <= bandwidth


def generate_input(batch: int, n: int, cond: int, seed: int, case: str = "dense") -> input_t:
    assert batch > 0, "batch must be positive"
    assert n > 0, "n must be positive"
    assert cond >= 0, "cond must be non-negative"

    device = "cuda" if torch.cuda.is_available() else "cpu"
    gen = torch.Generator(device=device)
    gen.manual_seed(seed)

    case = case.lower()
    a = torch.randn((batch, n, n), device=device, dtype=torch.float32, generator=gen)

    if case == "dense":
        a = _apply_column_scaling(a, cond)
    elif case == "upper":
        diag_boost = torch.linspace(1.0, 0.25, n, device=device, dtype=torch.float32)
        a = torch.triu(a)
        a.diagonal(dim1=-2, dim2=-1).add_(diag_boost)
        a = _apply_column_scaling(a, cond)
    elif case == "diagonal":
        diag = torch.randn((batch, n), device=device, dtype=torch.float32, generator=gen)
        diag = diag.sign().clamp(min=0.0).mul(2.0).sub(1.0) * torch.logspace(
            0.0, -float(max(cond, 2)), n, device=device, dtype=torch.float32
        )
        a = torch.diag_embed(diag)
    elif case == "rankdef":
        rank = max(1, (3 * n) // 4)
        a[:, :, rank:] = 0.0
        a = _apply_column_scaling(a, cond)
    elif case == "nearrank":
        rank = max(1, (3 * n) // 4)
        tail = n - rank
        if tail > 0:
            noise = torch.randn(
                (batch, n, tail), device=device, dtype=torch.float32, generator=gen
            )
            a[:, :, rank:] = a[:, :, :tail] + 1.0e-5 * noise
        a = _apply_column_scaling(a, cond)
    elif case == "clustered":
        scales = torch.ones((n,), device=device, dtype=torch.float32)
        scales[n // 2 :] = 4.0 * torch.finfo(torch.float32).eps
        if n >= 8:
            lo = max(0, n // 2 - 2)
            hi = min(n, n // 2 + 2)
            scales[lo:hi] = torch.sqrt(torch.tensor(torch.finfo(torch.float32).eps, device=device))
        a = a * scales
    elif case == "band":
        bandwidth = max(2, min(32, n // 32))
        a = a * _band_mask(n, bandwidth, device)
        diag_boost = torch.linspace(1.0, 0.5, n, device=device, dtype=torch.float32)
        a.diagonal(dim1=-2, dim2=-1).add_(diag_boost)
        a = _apply_column_scaling(a, cond)
    elif case == "nearcollinear":
        base = torch.randn((batch, n, 1), device=device, dtype=torch.float32, generator=gen)
        noise = torch.randn((batch, n, n), device=device, dtype=torch.float32, generator=gen)
        a = base.expand(batch, n, n) + 1.0e-4 * noise
        a = _apply_column_scaling(a, cond)
    elif case == "rowscale":
        row_cond = max(cond, 4)
        scales = torch.logspace(0.0, -float(row_cond), n, device=device, dtype=torch.float32)
        a = scales.reshape(1, n, 1) * a
    else:
        raise ValueError(f"unknown QR test case: {case}")

    return a.contiguous()


def ref_kernel(data: input_t) -> output_t:
    # Starter/reference path: correctness first; submissions compete on speed.
    return torch.geqrf(data)


def custom_kernel(data: input_t) -> output_t:
    h, tau = _qr_module.qr_householder_cuda(data)
    return h, tau


def _property_rtol(n: int, factor: float) -> float:
    eps = torch.finfo(torch.float32).eps
    return factor * max(n, 1) * eps


def _scaled_residual(
    residual: torch.Tensor,
    scale: torch.Tensor,
    n: int,
) -> torch.Tensor:
    eps = torch.finfo(torch.float32).eps
    return residual / (eps * max(n, 1) * scale.clamp_min(1e-30))


def _matrix_l1_norm(value: torch.Tensor) -> torch.Tensor:
    return torch.linalg.matrix_norm(value.double(), ord=1, dim=(-2, -1))


def _check_tensor(name: str, value: torch.Tensor, shape: tuple[int, ...], device: torch.device) -> str | None:
    if not isinstance(value, torch.Tensor):
        return f"{name} must be a torch.Tensor"
    if value.shape != shape:
        return f"{name} shape must be {shape}, got {tuple(value.shape)}"
    if value.dtype != torch.float32:
        return f"{name} dtype must be torch.float32, got {value.dtype}"
    if value.device != device:
        return f"{name} must be on {device}, got {value.device}"
    if not torch.isfinite(value).all().item():
        return f"{name} contains NaN or Inf"
    return None


def check_implementation(data: input_t, output: output_t) -> tuple[bool, str]:
    a = data
    batch, n, _ = a.shape
    factor_rtol = _property_rtol(n, _FACTOR_RTOL_FACTOR)
    orth_rtol = _property_rtol(n, _ORTH_RTOL_FACTOR)

    if not isinstance(output, tuple) or len(output) != 2:
        return False, "output must be a tuple `(H, tau)`"

    h, tau = output
    error = _check_tensor("H", h, (batch, n, n), a.device)
    if error is not None:
        return False, error
    error = _check_tensor("tau", tau, (batch, n), a.device)
    if error is not None:
        return False, error

    q = torch.linalg.householder_product(h, tau)
    r = torch.triu(h)
    a_check = a.double()
    q_check = q.double()
    r_check = r.double()
    projected = q_check.transpose(-1, -2) @ a_check
    factor_residual = _matrix_l1_norm(r_check - projected).amax()
    factor_scale = _matrix_l1_norm(a_check).amax()
    factor_allowed = factor_rtol * factor_scale
    factor_scaled = _scaled_residual(factor_residual, factor_scale, n)
    if factor_residual.item() > factor_allowed.item():
        return False, (
            "R - Q.T @ A is too large: "
            f"residual={factor_residual.item():.3g}, allowed={factor_allowed.item():.3g}, "
            f"scaled={factor_scaled.item():.3g}"
        )

    eye = torch.eye(n, device=a.device, dtype=torch.float64).expand(batch, n, n)
    qtq = q_check.transpose(-1, -2) @ q_check
    orth_residual = _matrix_l1_norm(qtq - eye).amax()
    orth_scale = _matrix_l1_norm(eye).amax()
    orth_allowed = orth_rtol * orth_scale
    orth_scaled = _scaled_residual(orth_residual, orth_scale, n)
    if orth_residual.item() > orth_allowed.item():
        return False, (
            "Q is not orthogonal enough: "
            f"residual={orth_residual.item():.3g}, allowed={orth_allowed.item():.3g}, "
            f"scaled={orth_scaled.item():.3g}"
        )

    lower = torch.tril(projected, diagonal=-1)
    tri_residual = _matrix_l1_norm(lower).amax()
    tri_scale = _matrix_l1_norm(a_check).amax()
    tri_scaled = _scaled_residual(tri_residual, tri_scale, n)

    recon = q_check @ r_check
    recon_residual = _matrix_l1_norm(recon - a_check).amax()
    recon_scale = _matrix_l1_norm(a_check).amax()
    recon_scaled = _scaled_residual(recon_residual, recon_scale, n)

    return True, (
        f"factor_rtol={factor_rtol:.3g}; "
        f"orth_rtol={orth_rtol:.3g}; "
        f"scaled_factor_residual={factor_scaled.item():.3g}; "
        f"scaled_reconstruction_residual={recon_scaled.item():.3g}; "
        f"scaled_triangular_residual={tri_scaled.item():.3g}; "
        f"scaled_orthogonality_residual={orth_scaled.item():.3g}; "
        f"batch={batch}; n={n}"
    )