submission 824419

#!POPCORN leaderboard qr_v2
#!POPCORN gpu B200
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

template <typename scalar_t>
__global__ void householder_qr_kernel(
    scalar_t* __restrict__ A, scalar_t* __restrict__ tau, int n)
{
    extern __shared__ unsigned char smem_raw[];
    scalar_t* v_shared = reinterpret_cast<scalar_t*>(smem_raw);
    scalar_t* red = v_shared + n;

    const int b = blockIdx.x;
    const int tid = threadIdx.x;
    const int nthreads = blockDim.x;
    scalar_t* Aglobal = A + (size_t)b * n * n;

    for (int k = 0; k < n; ++k) {
        const int rlen = n - k;

        for (int i = tid; i < rlen; i += nthreads)
            v_shared[i] = Aglobal[(size_t)(k + i) * n + k];
        __syncthreads();

        scalar_t local_max = (scalar_t)0;
        for (int i = tid; i < rlen; i += nthreads) {
            scalar_t av = fabsf((float)v_shared[i]);
            if (av > local_max) local_max = av;
        }
        red[tid] = local_max;
        __syncthreads();
        for (int s = nthreads / 2; s > 0; s >>= 1) {
            if (tid < s && red[tid + s] > red[tid]) red[tid] = red[tid + s];
            __syncthreads();
        }
        scalar_t colmax = red[0];
        __syncthreads();

        scalar_t norm;
        if (colmax == (scalar_t)0) {
            norm = (scalar_t)0;
        } else {
            scalar_t local_sumsq = (scalar_t)0;
            for (int i = tid; i < rlen; i += nthreads) {
                scalar_t vv = v_shared[i] / colmax;
                local_sumsq += vv * vv;
            }
            red[tid] = local_sumsq;
            __syncthreads();
            for (int s = nthreads / 2; s > 0; s >>= 1) {
                if (tid < s) red[tid] += red[tid + s];
                __syncthreads();
            }
            norm = colmax * sqrtf((float)red[0]);
            __syncthreads();
        }

        scalar_t x0 = v_shared[0];

        if (norm == (scalar_t)0) {
            if (tid == 0) tau[(size_t)b * n + k] = (scalar_t)0;
            __syncthreads();
            continue;
        }

        scalar_t sign_x0 = (x0 >= (scalar_t)0) ? (scalar_t)1 : (scalar_t)-1;
        scalar_t alpha = -sign_x0 * norm;
        scalar_t denom = x0 - alpha;
        scalar_t tau_k = (alpha - x0) / alpha;

        for (int i = tid + 1; i < rlen; i += nthreads)
            v_shared[i] = v_shared[i] / denom;
        __syncthreads();
        if (tid == 0) {
            v_shared[0] = (scalar_t)1;
            tau[(size_t)b * n + k] = tau_k;
        }
        __syncthreads();

        if (tid == 0) Aglobal[(size_t)k * n + k] = alpha;
        for (int i = tid + 1; i < rlen; i += nthreads)
            Aglobal[(size_t)(k + i) * n + k] = v_shared[i];
        __syncthreads();

        for (int j = k + 1 + tid; j < n; j += nthreads) {
            scalar_t dot = (scalar_t)0;
            for (int i = 0; i < rlen; ++i)
                dot += v_shared[i] * Aglobal[(size_t)(k + i) * n + j];
            scalar_t scale = tau_k * dot;
            for (int i = 0; i < rlen; ++i)
                Aglobal[(size_t)(k + i) * n + j] -= scale * v_shared[i];
        }
        __syncthreads();
    }
}

std::tuple<torch::Tensor, torch::Tensor> run_qr(torch::Tensor a) {
    TORCH_CHECK(a.is_cuda(), "input must be CUDA tensor");
    TORCH_CHECK(a.dim() == 3, "input must be batch x n x n");
    TORCH_CHECK(a.size(1) == a.size(2), "input must be square");
    TORCH_CHECK(a.scalar_type() == torch::kFloat32, "this kernel only supports fp32");

    auto A = a.contiguous().clone();
    int batch = A.size(0);
    int n = A.size(1);
    auto tau = torch::empty({batch, n}, A.options());

    int threads = 256;
    if (n < 256) {
        threads = 32;
        while (threads * 2 <= n && threads < 256) threads *= 2;
    }

    size_t smem_bytes = (size_t)n * sizeof(float) + (size_t)threads * sizeof(float);

    if (smem_bytes > 49152) {
        cudaFuncSetAttribute(
            householder_qr_kernel<float>,
            cudaFuncAttributeMaxDynamicSharedMemorySize,
            (int)smem_bytes);
    }
    TORCH_CHECK(smem_bytes <= 227 * 1024,
        "n=", n, " needs ", smem_bytes, " bytes of shared memory, exceeds budget");

    householder_qr_kernel<float><<<batch, threads, smem_bytes>>>(
        A.data_ptr<float>(), tau.data_ptr<float>(), n);

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) throw std::runtime_error(cudaGetErrorString(err));

    return std::make_tuple(A, tau);
}
"""

cpp_source = """
#include <torch/extension.h>
#include <tuple>
std::tuple<torch::Tensor, torch::Tensor> run_qr(torch::Tensor a);
"""

module = load_inline(
    name='custom_qr_module',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['run_qr'],
    verbose=True,
    extra_cuda_cflags=["-O3", "-use_fast_math"]
)

def custom_kernel(data: input_t) -> output_t:
    data = data.contiguous()
    h, tau = module.run_qr(data)
    return h, tau