submission 799674

#!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_SOURCE = r"""
#include <torch/extension.h>

#include <vector>

std::vector<torch::Tensor> qr_small(torch::Tensor input);
void configure_qr();
"""


_CUDA_SOURCE = r"""
#include <torch/extension.h>

#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

#include <cmath>
#include <vector>

namespace {

constexpr unsigned kFullMask = 0xffffffffu;

__device__ __forceinline__ float warp_sum(float value) {
#pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        value += __shfl_down_sync(kFullMask, value, offset);
    }
    return value;
}

template <int Threads>
__device__ __forceinline__ float block_sum(float value, float* warp_sums) {
    constexpr int Warps = Threads / 32;
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    value = warp_sum(value);
    if (lane == 0) {
        warp_sums[warp] = value;
    }
    __syncthreads();

    value = threadIdx.x < Warps ? warp_sums[lane] : 0.0f;
    if (warp == 0) {
        value = warp_sum(value);
    }
    if (threadIdx.x == 0) {
        warp_sums[0] = value;
    }
    __syncthreads();
    return warp_sums[0];
}

template <int N, int Threads>
__global__ void qr_shared_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    float* __restrict__ tau,
    int batch
) {
    constexpr int Warps = Threads / 32;
    const int matrix_id = blockIdx.x;
    if (matrix_id >= batch) {
        return;
    }

    extern __shared__ float shared[];
    float* matrix = shared;
    float* scratch = matrix + N * N;

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int matrix_offset = matrix_id * N * N;

    // Transpose once so each reflector column is contiguous across warp lanes.
    for (int index = tid; index < N * N; index += Threads) {
        const int row = index / N;
        const int col = index - row * N;
        matrix[col * N + row] = input[matrix_offset + index];
    }
    __syncthreads();

    for (int k = 0; k < N; ++k) {
        float local_norm = 0.0f;
        for (int row = k + 1 + tid; row < N; row += Threads) {
            const float value = matrix[k * N + row];
            local_norm = fmaf(value, value, local_norm);
        }
        const float norm_sq = block_sum<Threads>(local_norm, scratch);

        if (tid == 0) {
            const float alpha = matrix[k * N + k];
            const float xnorm = sqrtf(norm_sq);
            float tau_k = 0.0f;
            float scale = 0.0f;

            if (xnorm != 0.0f) {
                const float beta = -copysignf(hypotf(alpha, xnorm), alpha);
                tau_k = (beta - alpha) / beta;
                scale = 1.0f / (alpha - beta);
                matrix[k * N + k] = beta;
            }

            tau[matrix_id * N + k] = tau_k;
            scratch[0] = tau_k;
            scratch[1] = scale;
        }
        __syncthreads();

        const float tau_k = scratch[0];
        const float scale = scratch[1];
        if (tau_k != 0.0f) {
            for (int row = k + 1 + tid; row < N; row += Threads) {
                matrix[k * N + row] *= scale;
            }
        }
        __syncthreads();

        if (tau_k != 0.0f) {
            for (int col = k + 1 + warp; col < N; col += Warps) {
                float dot = 0.0f;
                for (int row = k + lane; row < N; row += 32) {
                    const float v = row == k ? 1.0f : matrix[k * N + row];
                    dot = fmaf(v, matrix[col * N + row], dot);
                }
                dot = warp_sum(dot);
                const float weight = __shfl_sync(kFullMask, dot, 0) * tau_k;

                for (int row = k + lane; row < N; row += 32) {
                    const float v = row == k ? 1.0f : matrix[k * N + row];
                    matrix[col * N + row] =
                        fmaf(-v, weight, matrix[col * N + row]);
                }
            }
        }
        __syncthreads();
    }

    for (int index = tid; index < N * N; index += Threads) {
        const int row = index / N;
        const int col = index - row * N;
        output[matrix_offset + index] = matrix[col * N + row];
    }
}

template <int N, int Threads>
void launch_qr(
    const torch::Tensor& input,
    torch::Tensor& output,
    torch::Tensor& tau
) {
    constexpr int SharedFloats = N * N + Threads / 32 + 2;
    const int batch = static_cast<int>(input.size(0));
    qr_shared_kernel<N, Threads>
        <<<batch, Threads, SharedFloats * sizeof(float)>>>(
            input.data_ptr<float>(),
            output.data_ptr<float>(),
            tau.data_ptr<float>(),
            batch
        );
    const cudaError_t error = cudaGetLastError();
    TORCH_CHECK(
        error == cudaSuccess,
        "qr_shared_kernel launch failed: ",
        cudaGetErrorString(error)
    );
}

}  // namespace

void configure_qr() {
    constexpr int SharedBytes176 = (176 * 176 + 1024 / 32 + 2) * sizeof(float);
    const cudaError_t error = cudaFuncSetAttribute(
        qr_shared_kernel<176, 1024>,
        cudaFuncAttributeMaxDynamicSharedMemorySize,
        SharedBytes176
    );
    TORCH_CHECK(
        error == cudaSuccess,
        "failed to configure qr_shared_kernel: ",
        cudaGetErrorString(error)
    );
}

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

    const int64_t n = input.size(1);
    TORCH_CHECK(n == 32 || n == 176, "native path supports n=32 or n=176");

    c10::cuda::CUDAGuard device_guard(input.device());
    auto output = torch::empty_like(input);
    auto tau = torch::empty({input.size(0), n}, input.options());

    if (n == 32) {
        launch_qr<32, 512>(input, output, tau);
    } else {
        launch_qr<176, 1024>(input, output, tau);
    }
    return {output, tau};
}
"""


_qr_extension = load_inline(
    name="qr1_shared_householder_v3",
    cpp_sources=[_CPP_SOURCE],
    cuda_sources=[_CUDA_SOURCE],
    functions=["qr_small", "configure_qr"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)
_qr_extension.configure_qr()


def custom_kernel(data: input_t) -> output_t:
    n = data.shape[-1]
    if n == 32 or n == 176:
        h, tau = _qr_extension.qr_small(data)
        return h, tau
    if n >= 352:
        rank = (3 * n) // 4
        tail_nonzero = torch.count_nonzero(data[:, :, rank:], dim=(1, 2))
        zero_mask = tail_nonzero == 0
        zero_count = int(torch.count_nonzero(zero_mask).item())
        if zero_count == data.shape[0]:
            h_head, tau_head = torch.geqrf(data[:, :, :rank].contiguous())
            h = torch.empty_like(data)
            h[:, :, :rank] = h_head
            h[:, :, rank:].zero_()
            tau = torch.empty((data.shape[0], n), device=data.device, dtype=data.dtype)
            tau[:, :rank] = tau_head
            tau[:, rank:].zero_()
            return h, tau
        if zero_count > 0:
            zero_idx = torch.nonzero(zero_mask, as_tuple=False).flatten()
            full_idx = torch.nonzero(~zero_mask, as_tuple=False).flatten()

            h = torch.empty_like(data)
            tau = torch.empty((data.shape[0], n), device=data.device, dtype=data.dtype)

            full_data = data.index_select(0, full_idx).contiguous()
            h_full, tau_full = torch.geqrf(full_data)
            h.index_copy_(0, full_idx, h_full)
            tau.index_copy_(0, full_idx, tau_full)

            zero_data = data.index_select(0, zero_idx)[:, :, :rank].contiguous()
            h_head, tau_head = torch.geqrf(zero_data)
            h_zero = torch.empty((zero_count, n, n), device=data.device, dtype=data.dtype)
            h_zero[:, :, :rank] = h_head
            h_zero[:, :, rank:].zero_()
            tau_zero = torch.empty((zero_count, n), device=data.device, dtype=data.dtype)
            tau_zero[:, :rank] = tau_head
            tau_zero[:, rank:].zero_()
            h.index_copy_(0, zero_idx, h_zero)
            tau.index_copy_(0, zero_idx, tau_zero)
            return h, tau
    return torch.geqrf(data)