submission 797694

#!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 = r"""
void transpose_square_inplace(torch::Tensor A);
void factor_col_tlayout(torch::Tensor A, torch::Tensor tau, int64_t k);
void apply_reflector_tlayout(torch::Tensor A, torch::Tensor tau, int64_t k);
"""

CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <stdexcept>

#define BLOCK 256

__global__ void transpose_square_inplace_kernel(
    float* __restrict__ A,
    int batch,
    int n
) {
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)batch * n * n;
    if (tid >= total) return;

    long long mat_size = (long long)n * n;
    int b = (int)(tid / mat_size);
    long long p = tid - (long long)b * mat_size;
    int i = (int)(p / n);
    int j = (int)(p - (long long)i * n);

    if (i < j) {
        long long base = (long long)b * mat_size;
        long long idx1 = base + (long long)i * n + j;
        long long idx2 = base + (long long)j * n + i;
        float tmp = A[idx1];
        A[idx1] = A[idx2];
        A[idx2] = tmp;
    }
}

// Internal layout after transpose:
// physical A[phys_row, phys_col] = logical A[phys_col, phys_row].
// Therefore logical A[row, col] lives at physical A[col, row].
// Factor logical column k by reading/writing physical row k, columns k:n.
__global__ void factor_col_tlayout_kernel(
    float* __restrict__ A,
    float* __restrict__ tau,
    int batch,
    int n,
    int k
) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int m = n - k;

    __shared__ float sh_sum[BLOCK];
    __shared__ float sh_denom;
    __shared__ int sh_general;

    long long base = (long long)b * n * n;
    long long row_base = base + (long long)k * n;

    float local_sum = 0.0f;

    // sigma = sum_{r=1}^{m-1} logical A[k+r, k]^2
    // logical A[k+r, k] == physical A[k, k+r]
    for (int r = tid + 1; r < m; r += BLOCK) {
        float x = A[row_base + k + r];
        local_sum += x * x;
    }

    sh_sum[tid] = local_sum;
    __syncthreads();

    for (int stride = BLOCK / 2; stride > 0; stride >>= 1) {
        if (tid < stride) sh_sum[tid] += sh_sum[tid + stride];
        __syncthreads();
    }

    float sigma = sh_sum[0];

    if (tid == 0) {
        float alpha = A[row_base + k];
        float beta, tau_k, denom;
        int general;

        if (sigma == 0.0f && alpha >= 0.0f) {
            beta = alpha;
            tau_k = 0.0f;
            denom = 1.0f;
            general = 0;
        } else if (sigma == 0.0f && alpha < 0.0f) {
            beta = -alpha;
            tau_k = 2.0f;
            denom = 1.0f;
            general = 0;
        } else {
            float norm = sqrtf(alpha * alpha + sigma);
            float s = alpha >= 0.0f ? 1.0f : -1.0f;
            beta = -s * norm;
            denom = alpha - beta;
            tau_k = (beta - alpha) / beta;
            general = 1;
        }

        A[row_base + k] = beta;
        tau[(long long)b * n + k] = tau_k;
        sh_denom = denom;
        sh_general = general;
    }

    __syncthreads();

    // Store v_tail in logical A[k+1:n, k], i.e. physical A[k, k+1:n].
    for (int r = tid + 1; r < m; r += BLOCK) {
        long long idx = row_base + k + r;
        float x = A[idx];
        A[idx] = sh_general ? x / sh_denom : 0.0f;
    }
}

// Apply reflector to logical trailing columns j > k.
// One block owns one logical trailing column j for one batch item.
// In physical layout this is row j, columns k:n, so the row sweep is contiguous.
__global__ void apply_reflector_tlayout_kernel(
    float* __restrict__ A,
    const float* __restrict__ tau,
    int batch,
    int n,
    int k
) {
    int b = blockIdx.x;
    int col_offset = blockIdx.y;
    int j = k + 1 + col_offset;
    if (j >= n) return;

    int tid = threadIdx.x;
    int m = n - k;
    long long base = (long long)b * n * n;

    long long v_base = base + (long long)k * n;
    long long a_base = base + (long long)j * n;

    __shared__ float sh_sum[BLOCK];
    float tau_k = tau[(long long)b * n + k];

    // dot = v^T logical A[k:n, j]
    // logical A[k+r, j] == physical A[j, k+r]
    float local_sum = 0.0f;
    for (int r = tid; r < m; r += BLOCK) {
        float v = (r == 0) ? 1.0f : A[v_base + k + r];
        float a = A[a_base + k + r];
        local_sum += v * a;
    }

    sh_sum[tid] = local_sum;
    __syncthreads();

    for (int stride = BLOCK / 2; stride > 0; stride >>= 1) {
        if (tid < stride) sh_sum[tid] += sh_sum[tid + stride];
        __syncthreads();
    }

    float scaled_dot = tau_k * sh_sum[0];

    // logical A[k:n, j] -= v * scaled_dot
    for (int r = tid; r < m; r += BLOCK) {
        float v = (r == 0) ? 1.0f : A[v_base + k + r];
        long long idx = a_base + k + r;
        A[idx] -= v * scaled_dot;
    }
}

void transpose_square_inplace(torch::Tensor A) {
    TORCH_CHECK(A.is_cuda(), "A must be CUDA");
    TORCH_CHECK(A.dtype() == torch::kFloat32, "A must be float32");
    TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
    TORCH_CHECK(A.dim() == 3, "A must have shape [batch, n, n]");

    int batch = A.size(0);
    int n = A.size(1);
    TORCH_CHECK(A.size(2) == n, "A must be square");

    long long total = (long long)batch * n * n;
    int blocks = (int)((total + BLOCK - 1) / BLOCK);

    transpose_square_inplace_kernel<<<blocks, BLOCK>>>(
        A.data_ptr<float>(), batch, n
    );

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

void factor_col_tlayout(torch::Tensor A, torch::Tensor tau, int64_t k) {
    TORCH_CHECK(A.is_cuda(), "A must be CUDA");
    TORCH_CHECK(tau.is_cuda(), "tau must be CUDA");
    TORCH_CHECK(A.dtype() == torch::kFloat32, "A must be float32");
    TORCH_CHECK(tau.dtype() == torch::kFloat32, "tau must be float32");
    TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
    TORCH_CHECK(tau.is_contiguous(), "tau must be contiguous");

    int batch = A.size(0);
    int n = A.size(1);

    factor_col_tlayout_kernel<<<dim3(batch), dim3(BLOCK)>>>(
        A.data_ptr<float>(), tau.data_ptr<float>(), batch, n, (int)k
    );

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

void apply_reflector_tlayout(torch::Tensor A, torch::Tensor tau, int64_t k) {
    TORCH_CHECK(A.is_cuda(), "A must be CUDA");
    TORCH_CHECK(tau.is_cuda(), "tau must be CUDA");
    TORCH_CHECK(A.dtype() == torch::kFloat32, "A must be float32");
    TORCH_CHECK(tau.dtype() == torch::kFloat32, "tau must be float32");
    TORCH_CHECK(A.is_contiguous(), "A must be contiguous");
    TORCH_CHECK(tau.is_contiguous(), "tau must be contiguous");

    int batch = A.size(0);
    int n = A.size(1);
    if (k + 1 >= n) return;

    apply_reflector_tlayout_kernel<<<dim3(batch, n - (int)k - 1), dim3(BLOCK)>>>(
        A.data_ptr<float>(), tau.data_ptr<float>(), batch, n, (int)k
    );

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

module = load_inline(
    name="qr_full_cuda_tlayout_strict",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["transpose_square_inplace", "factor_col_tlayout", "apply_reflector_tlayout"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    batch, n, _ = data.shape
    if n > 2048:
        return torch.geqrf(data)
    A = data.contiguous().clone()
    tau = torch.empty((batch, n), device=A.device, dtype=A.dtype)

    # Convert physical layout to A^T. Kernels below still factor the original
    # logical matrix, but index it through transposed physical storage.
    module.transpose_square_inplace(A)

    for k in range(n):
        module.factor_col_tlayout(A, tau, k)
        module.apply_reflector_tlayout(A, tau, k)

    # Convert compact Householder H back to the expected physical layout.
    module.transpose_square_inplace(A)

    return A, tau