submission 896090

#!POPCORN leaderboard cholesky
#!POPCORN gpu B200

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

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

// ==========================================================================
// N=32: Warp-level register kernel with coalesced I/O via shared memory.
// ==========================================================================
__global__ void __launch_bounds__(32)
chol32(const float* __restrict__ A, float* __restrict__ L, int batch)
{
    const int bid = blockIdx.x;
    if (bid >= batch) return;
    const int lane = threadIdx.x;  // 0..31

    const float* __restrict__ Ap = A + bid * 1024;
    float* __restrict__ Lp = L + bid * 1024;

    __shared__ float sm[32][33];
    #pragma unroll
    for (int i = 0; i < 32; i++) {
        sm[i][lane] = __ldg(&Ap[i * 32 + lane]);
    }
    __syncwarp();

    float r[32];
    #pragma unroll
    for (int j = 0; j < 32; j++) {
        r[j] = (j <= lane) ? sm[lane][j] : 0.0f;
    }

    #pragma unroll
    for (int j = 0; j < 32; j++) {
        float d = 0.0f;
        if (lane == j) {
            float s = 0.0f;
            #pragma unroll
            for (int k = 0; k < j; k++) s += r[k] * r[k];
            d = sqrtf(fmaxf(r[j] - s, 0.0f));
            d = fmaxf(d, 1e-20f);
            r[j] = d;
        }
        d = __shfl_sync(0xFFFFFFFF, d, j);
        float inv_d = __fdividef(1.0f, d);

        float dot = 0.0f;
        #pragma unroll
        for (int k = 0; k < j; k++) {
            float ljk = __shfl_sync(0xFFFFFFFF, r[k], j);
            dot += r[k] * ljk;
        }
        if (lane > j) r[j] = (r[j] - dot) * inv_d;
    }

    #pragma unroll
    for (int j = 0; j < 32; j++) {
        sm[lane][j] = (j <= lane) ? r[j] : 0.0f;
    }
    __syncwarp();

    #pragma unroll
    for (int i = 0; i < 32; i++) {
        Lp[i * 32 + lane] = sm[i][lane];
    }
}

// ==========================================================================
// N=64/128: Shared-memory kernel. One thread per row.
// ==========================================================================
template <int N>
__global__ void __launch_bounds__(N)
cholSM(const float* __restrict__ A, float* __restrict__ L, int batch)
{
    const int bid = blockIdx.x;
    if (bid >= batch) return;
    const int tid = threadIdx.x;
    constexpr int S = N + 1;

    extern __shared__ float sm[];

    const float* __restrict__ Ap = A + (long long)bid * N * N;
    float* __restrict__ Lp = L + (long long)bid * N * N;

    #pragma unroll 4
    for (int i = tid; i < N * N; i += N) {
        int row = i >> __builtin_ctz(N);
        int col = i & (N - 1);
        sm[row * S + col] = (col <= row) ? __ldg(&Ap[i]) : 0.0f;
    }
    __syncthreads();

    for (int j = 0; j < N; j++) {
        if (tid == j) {
            float s = 0.0f;
            for (int k = 0; k < j; k++) {
                float v = sm[j * S + k];
                s += v * v;
            }
            float diag = sm[j * S + j] - s;
            sm[j * S + j] = sqrtf(fmaxf(diag, 0.0f)) + 1e-30f;
        }
        __syncthreads();

        if (tid > j) {
            float s = 0.0f;
            float ljj_inv = __fdividef(1.0f, sm[j * S + j]);
            for (int k = 0; k < j; k++) {
                s += sm[tid * S + k] * sm[j * S + k];
            }
            sm[tid * S + j] = (sm[tid * S + j] - s) * ljj_inv;
        }
        __syncthreads();
    }

    #pragma unroll 4
    for (int i = tid; i < N * N; i += N) {
        int row = i >> __builtin_ctz(N);
        int col = i & (N - 1);
        Lp[i] = (col <= row) ? sm[row * S + col] : 0.0f;
    }
}

// ==========================================================================
// Launcher
// ==========================================================================
torch::Tensor cholesky_small(torch::Tensor A)
{
    TORCH_CHECK(A.is_cuda() && A.is_contiguous() && A.scalar_type() == at::kFloat);
    const int batch = A.size(0);
    const int n     = A.size(1);
    auto L = torch::empty_like(A);
    const float* ap = A.data_ptr<float>();
    float*       lp = L.data_ptr<float>();

    if (n == 32) {
        chol32<<<batch, 32>>>(ap, lp, batch);
    } else if (n == 64) {
        constexpr int smem = 64 * 65 * sizeof(float);
        cholSM<64><<<batch, 64, smem>>>(ap, lp, batch);
    } else if (n == 128) {
        constexpr int smem = 128 * 129 * sizeof(float);
        cudaFuncSetAttribute(cholSM<128>,
            cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
        cholSM<128><<<batch, 128, smem>>>(ap, lp, batch);
    } else {
        TORCH_CHECK(false, "cholesky_small: unsupported n=", n);
    }
    return L;
}
"""

_ext = load_inline(
    name="chol_ext_v2",
    cpp_sources=["torch::Tensor cholesky_small(torch::Tensor A);"],
    cuda_sources=[_CUDA_SRC],
    functions=["cholesky_small"],
    extra_cuda_cflags=[
        "-O3",
        "--use_fast_math",
        "-lineinfo",
        "--ptxas-options=-v",
        "-gencode=arch=compute_100,code=sm_100",  # B200 Blackwell
        "-gencode=arch=compute_90,code=sm_90",    # H100 fallback
    ],
    verbose=False,
)

def custom_kernel(data: input_t) -> output_t:
    """
    Batched dense Cholesky factorization optimized for B200.
    
    Strategy:
    - n ≤ 128: Custom CUDA kernels
    - n ≥ 256: cuSOLVER via torch.linalg.cholesky
    """
    A = data.contiguous()
    n = A.size(1)
    if n <= 128:
        return _ext.cholesky_small(A)
    return torch.linalg.cholesky(A, upper=False)