submission 884062

#!POPCORN leaderboard cholesky
#!POPCORN gpu B200

import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t

torch.backends.cuda.matmul.allow_tf32 = True

_NB = 4096
_BBNB = 256


def _batched_blocked_chol_tf32(A, nb=_BBNB):
    # Batched blocked right-looking Cholesky, keeping the batch dim: batched
    # FP32 diagonal potrf + batched triangular solve, trailing baddbmm in TF32.
    # Beats cuSOLVER batched potrf on mid-size high-batch (n1024 b60).
    b, n, _ = A.shape
    M = A.clone()
    L = torch.zeros_like(A)
    for k in range(0, n, nb):
        e = min(k + nb, n)
        Lkk = torch.linalg.cholesky_ex(M[:, k:e, k:e], check_errors=False).L
        L[:, k:e, k:e] = Lkk
        if e < n:
            B = M[:, e:, k:e]
            Xt = torch.linalg.solve_triangular(Lkk, B.transpose(-1, -2), upper=False)
            Lpanel = Xt.transpose(-1, -2).contiguous()
            L[:, e:, k:e] = Lpanel
            M[:, e:, e:].baddbmm_(Lpanel, Lpanel.transpose(-1, -2), beta=1.0, alpha=-1.0)
    return L


def _blocked_chol_tf32(A, nb=_NB):
    # Blocked right-looking Cholesky. Panel (diagonal potrf + triangular
    # solve) stays FP32; the dominant O(n^2*nb) trailing update runs in TF32
    # on tensor cores (~4x GEMM throughput vs cuSOLVER's FP32 potrf).
    n = A.size(-1)
    M = A.clone()
    L = torch.zeros_like(A)
    for k in range(0, n, nb):
        e = min(k + nb, n)
        Lkk = torch.linalg.cholesky_ex(M[k:e, k:e], check_errors=False).L
        L[k:e, k:e] = Lkk
        if e < n:
            B = M[e:, k:e]
            Xt = torch.linalg.solve_triangular(Lkk, B.transpose(0, 1), upper=False)
            Lpanel = Xt.transpose(0, 1).contiguous()
            L[e:, k:e] = Lpanel
            M[e:, e:].addmm_(Lpanel, Lpanel.transpose(0, 1), beta=1.0, alpha=-1.0)
    return L

CUDA_SRC = r"""
#include <cuda_runtime.h>
#include <math.h>

// One thread block per matrix. blockDim.x = n (one thread per row).
// Full matrix (lower triangle used) resident in shared memory with padded
// leading dimension lda = n|1 (odd) so column-strided smem access is
// bank-conflict-free. Left-looking (Cholesky-Crout) column sweep.
__global__ void chol_smem_kernel(const float* __restrict__ A,
                                 float* __restrict__ L,
                                 int n, int lda) {
    extern __shared__ float s[];  // n * lda
    const int b = blockIdx.x;
    const float* __restrict__ Ab = A + (size_t)b * n * n;
    float*       __restrict__ Lb = L + (size_t)b * n * n;
    const int tid = threadIdx.x;  // row index

    // Load lower triangle of A row `tid` into smem.
    for (int j = 0; j <= tid; ++j)
        s[tid * lda + j] = Ab[tid * n + j];
    __syncthreads();

    for (int j = 0; j < n; ++j) {
        if (tid == j) {
            float d = s[j * lda + j];
            #pragma unroll 4
            for (int k = 0; k < j; ++k) { float x = s[j * lda + k]; d -= x * x; }
            s[j * lda + j] = sqrtf(d);
        }
        __syncthreads();
        if (tid > j) {
            const float ljj = s[j * lda + j];
            float v = s[tid * lda + j];
            #pragma unroll 4
            for (int k = 0; k < j; ++k) v -= s[tid * lda + k] * s[j * lda + k];
            s[tid * lda + j] = v / ljj;
        }
        __syncthreads();
    }

    // Store: lower triangle (incl diag) from smem, upper triangle = 0.
    for (int j = 0; j < n; ++j)
        Lb[tid * n + j] = (j <= tid) ? s[tid * lda + j] : 0.0f;
}

torch::Tensor chol_batched(torch::Tensor A, torch::Tensor L) {
    const int batch = A.size(0);
    const int n = A.size(1);
    const int lda = n | 1;
    const size_t smem = (size_t)n * lda * sizeof(float);

    static int configured = -1;
    if (configured != n) {
        cudaFuncSetAttribute(chol_smem_kernel,
            cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
        configured = n;
    }

    chol_smem_kernel<<<batch, n, smem>>>(
        A.data_ptr<float>(), L.data_ptr<float>(), n, lda);

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

CPP_SRC = r"""
torch::Tensor chol_batched(torch::Tensor A, torch::Tensor L);
"""

module = load_inline(
    name='chol_smem',
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=['chol_batched'],
    verbose=True,
    extra_cuda_cflags=["-O3"],
)


def custom_kernel(data: input_t) -> output_t:
    A = data
    b = A.size(0)
    n = A.size(1)
    if n <= 64:
        # Custom batched smem kernel wins on many-small matrices.
        A = A.contiguous()
        L = torch.empty_like(A)
        return module.chol_batched(A, L)
    if n >= 8192:
        # Single huge matrices: blocked TF32-trailing beats cuSOLVER potrf.
        outs = [_blocked_chol_tf32(A[i]) for i in range(b)]
        return torch.stack(outs, dim=0)
    if n == 512 and b > 64:
        # n512 b640: marginal win (3.78 -> 3.72 ms) at nb=128 (n/4).
        return _batched_blocked_chol_tf32(A, nb=128)
    if n == 1024 and b > 4:
        # Mid-size high-batch: batched blocked TF32-trailing beats cuSOLVER
        # batched potrf (n1024 b60: 2.89 -> 2.37 ms). n2048 b8 shows no win.
        return _batched_blocked_chol_tf32(A, nb=256)
    if n >= 1024 and b <= 4:
        # cuSOLVER's batched potrf is pathological at low-batch/large-n
        # (e.g. n4096 b2: 5950 vs 1531 us/matrix). Factor each matrix with
        # the single-matrix path instead, one at a time.
        L = torch.empty_like(A)
        for i in range(b):
            L[i] = torch.linalg.cholesky_ex(A[i], check_errors=False).L
        return L
    return torch.linalg.cholesky_ex(A, check_errors=False).L