submission 885163

#!POPCORN leaderboard cholesky
#!POPCORN gpu B200
"""Batched dense Cholesky, GPU Mode `cholesky` (B200).

s2 = s1 + a custom BATCHED left-looking Cholesky for medium n (default n==512,
the heaviest shape n=512 b=640). Profiling showed cuSOLVER's batched path uses
block-16 CUDA-core kernels (NO tensor cores) + a full input clone (~1.19ms at
n=512 b640). This kernel:
  * left-looking (reads A panel-by-panel, writes L) → no full clone.
  * FLOP-heavy trailing GEMM  W = A_panel - L_left @ L_top^T  on cuBLAS
    strided-batched TF32 tensor cores.
  * small nb x nb diagonal factor + TRSM in one on-chip FP32 panel kernel.

s1 tiers kept: n==32 warp-reg, n==64 block-reg, n>=1024 & b<=8 per-matrix
non-batched potrf reroute, else cuSOLVER batched.
"""

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

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

_NB = 32  # panel width for the batched left-looking path

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>
#include <cmath>

// ================= s1 tiny-n kernels =================
template<int N>
__global__ void chol_warp_kernel(const float* __restrict__ A,
                                 float* __restrict__ L, int batch) {
    const int wpb = blockDim.x >> 5;
    const int m = blockIdx.x * wpb + (threadIdx.x >> 5);
    if (m >= batch) return;
    const int lane = threadIdx.x & 31;
    const float* Ab = A + (size_t)m * N * N;
    float* Lb = L + (size_t)m * N * N;
    float a[N];
    #pragma unroll
    for (int c = 0; c < N; ++c) a[c] = (c <= lane) ? Ab[lane * N + c] : 0.0f;
    #pragma unroll
    for (int j = 0; j < N; ++j) {
        float d = sqrtf(__shfl_sync(0xffffffffu, a[j], j));
        float linv = 1.0f / d;
        float lij = (lane > j) ? a[j] * linv : (lane == j ? d : 0.0f);
        if (lane == j) a[j] = d; else if (lane > j) a[j] = lij;
        #pragma unroll
        for (int k = j + 1; k < N; ++k) {
            float lkj = __shfl_sync(0xffffffffu, lij, k);
            if (lane >= k) a[k] -= lij * lkj;
        }
    }
    #pragma unroll
    for (int c = 0; c < N; ++c) Lb[lane * N + c] = (c <= lane) ? a[c] : 0.0f;
}

template<int N>
__global__ void chol_block_kernel(const float* __restrict__ A,
                                   float* __restrict__ L) {
    __shared__ float col[N];
    const int b = blockIdx.x;
    const int i = threadIdx.x;
    const float* Ab = A + (size_t)b * N * N;
    float* Lb = L + (size_t)b * N * N;
    float a[N];
    #pragma unroll
    for (int c = 0; c < N; ++c) a[c] = (c <= i) ? Ab[i * N + c] : 0.0f;
    #pragma unroll
    for (int j = 0; j < N; ++j) {
        if (i == j) col[j] = a[j];
        __syncthreads();
        float d = sqrtf(col[j]);
        float linv = 1.0f / d;
        float lij = (i > j) ? a[j] * linv : (i == j ? d : 0.0f);
        __syncthreads();
        col[i] = lij;
        __syncthreads();
        float ci = col[i];
        #pragma unroll
        for (int k = j + 1; k < N; ++k) if (i >= k) a[k] -= ci * col[k];
        if (i == j) a[j] = d; else if (i > j) a[j] = lij;
        __syncthreads();
    }
    #pragma unroll
    for (int c = 0; c < N; ++c) Lb[i * N + c] = (c <= i) ? a[c] : 0.0f;
}

torch::Tensor chol_warp(torch::Tensor A) {
    const int batch = (int)A.size(0);
    auto L = torch::empty_like(A);
    const int wpb = 8, threads = wpb * 32, blocks = (batch + wpb - 1) / wpb;
    chol_warp_kernel<32><<<blocks, threads>>>(A.data_ptr<float>(), L.data_ptr<float>(), batch);
    return L;
}
torch::Tensor chol_block(torch::Tensor A) {
    const int batch = (int)A.size(0), n = (int)A.size(2);
    auto L = torch::empty_like(A);
    auto pa = A.data_ptr<float>(); auto pl = L.data_ptr<float>();
    if (n == 64) chol_block_kernel<64><<<batch, 64>>>(pa, pl);
    else throw std::runtime_error("chol_block unsupported n");
    return L;
}

// ================= batched left-looking (medium n) =================
static cublasHandle_t cbHandle() {
    static cublasHandle_t h = nullptr;
    if (!h) { if (cublasCreate(&h) != CUBLAS_STATUS_SUCCESS) throw std::runtime_error("cublasCreate"); }
    return h;
}

// Row-major C(m x n) = alpha*opA(A)(m x k) @ opB(B)(k x n) + beta*C, strided-batched.
static void gemm_rm(bool tA, bool tB, int m, int n, int k, float alpha,
                    const float* A, int lda, long long sA,
                    const float* B, int ldb, long long sB, float beta,
                    float* C, int ldc, long long sC, int batch, cublasComputeType_t ct) {
    cublasOperation_t opA = tA ? CUBLAS_OP_T : CUBLAS_OP_N;
    cublasOperation_t opB = tB ? CUBLAS_OP_T : CUBLAS_OP_N;
    cublasStatus_t st = cublasGemmStridedBatchedEx(
        cbHandle(), opB, opA, n, m, k, &alpha,
        B, CUDA_R_32F, ldb, sB, A, CUDA_R_32F, lda, sA,
        &beta, C, CUDA_R_32F, ldc, sC, batch, ct, CUBLAS_GEMM_DEFAULT);
    if (st != CUBLAS_STATUS_SUCCESS) throw std::runtime_error("gemm " + std::to_string((int)st));
}

// Copy A[k:k+R, k:k+cur] -> L[same]  (panel init for the left-looking GEMM).
__global__ void copy_panel_kernel(const float* __restrict__ A, float* __restrict__ L,
                                  int n, int k, int cur, int R) {
    const int b = blockIdx.z;
    const float* Ab = A + (size_t)b * n * n;
    float* Lb = L + (size_t)b * n * n;
    const long long tot = (long long)R * cur;
    for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < tot; idx += (long long)gridDim.x * blockDim.x) {
        int r = idx / cur, c = idx % cur;
        Lb[(size_t)(k + r) * n + (k + c)] = Ab[(size_t)(k + r) * n + (k + c)];
    }
}

// Fast NB=32 panel: warp0 factors the 32x32 diagonal block in registers via
// warp-shuffle (NO __syncthreads chain), then all threads do the TRSM below.
__global__ void panel_factor_warp32(float* __restrict__ L, int n, int k, int R) {
    __shared__ float D[32][33];   // [c][p] = L11[c][p]; +1 pad avoids bank conflicts
    const int b = blockIdx.x;
    float* Lb = L + (size_t)b * n * n;
    const int tid = threadIdx.x, nt = blockDim.x;
    const int warp = tid >> 5, lane = tid & 31;
    if (warp == 0) {
        float a[32];
        #pragma unroll
        for (int c = 0; c < 32; ++c)
            a[c] = (c <= lane) ? Lb[(size_t)(k + lane) * n + (k + c)] : 0.0f;
        #pragma unroll
        for (int j = 0; j < 32; ++j) {
            float d = sqrtf(__shfl_sync(0xffffffffu, a[j], j));
            float linv = 1.0f / d;
            float lij = (lane > j) ? a[j] * linv : (lane == j ? d : 0.0f);
            if (lane == j) a[j] = d; else if (lane > j) a[j] = lij;
            #pragma unroll
            for (int kk = j + 1; kk < 32; ++kk) {
                float lkj = __shfl_sync(0xffffffffu, lij, kk);
                if (lane >= kk) a[kk] -= lij * lkj;
            }
        }
        #pragma unroll
        for (int c = 0; c < 32; ++c) {
            float v = (c <= lane) ? a[c] : 0.0f;
            D[lane][c] = v;                              // D[row][col]
            Lb[(size_t)(k + lane) * n + (k + c)] = v;
        }
    }
    __syncthreads();
    for (int i = 32 + tid; i < R; i += nt) {
        float row[32];
        #pragma unroll
        for (int c = 0; c < 32; ++c) {
            float s = Lb[(size_t)(k + i) * n + (k + c)];
            #pragma unroll
            for (int p = 0; p < c; ++p) s -= row[p] * D[c][p];
            row[c] = s / D[c][c];
            Lb[(size_t)(k + i) * n + (k + c)] = row[c];
        }
    }
}

// Factor the panel L[k:k+R, k:k+cur] (already = updated W) in place:
// diagonal cur x cur Cholesky (FP32) + TRSM of the rows below. One block/matrix.
template<int NB>
__global__ void panel_factor_kernel(float* __restrict__ L, int n, int k, int R) {
    __shared__ float D[NB][NB];
    const int b = blockIdx.x;
    float* Lb = L + (size_t)b * n * n;
    const int tid = threadIdx.x, nt = blockDim.x;
    // load diagonal block (lower); zero strict upper
    for (int idx = tid; idx < NB * NB; idx += nt) {
        int r = idx / NB, c = idx % NB;
        D[r][c] = (r >= c) ? Lb[(size_t)(k + r) * n + (k + c)] : 0.0f;
    }
    __syncthreads();
    // left-looking factor of the diagonal block
    for (int j = 0; j < NB; ++j) {
        if (tid == 0) {
            float s = D[j][j];
            for (int p = 0; p < j; ++p) s -= D[j][p] * D[j][p];
            D[j][j] = sqrtf(s);
        }
        __syncthreads();
        float djj = D[j][j];
        for (int i = j + 1 + tid; i < NB; i += nt) {
            float s = D[i][j];
            for (int p = 0; p < j; ++p) s -= D[i][p] * D[j][p];
            D[i][j] = s / djj;
        }
        __syncthreads();
    }
    // write diagonal block back (lower kept, strict upper zeroed)
    for (int idx = tid; idx < NB * NB; idx += nt) {
        int r = idx / NB, c = idx % NB;
        Lb[(size_t)(k + r) * n + (k + c)] = (r >= c) ? D[r][c] : 0.0f;
    }
    // TRSM: rows i in [NB, R): L21[i,c] = (W[i,c] - sum_{p<c} L21[i,p]*D[c,p]) / D[c,c]
    for (int i = NB + tid; i < R; i += nt) {
        float row[NB];
        #pragma unroll
        for (int c = 0; c < NB; ++c) {
            float s = Lb[(size_t)(k + i) * n + (k + c)];
            #pragma unroll
            for (int p = 0; p < c; ++p) s -= row[p] * D[c][p];
            row[c] = s / D[c][c];
            Lb[(size_t)(k + i) * n + (k + c)] = row[c];
        }
    }
}

torch::Tensor chol_batched(torch::Tensor A, int64_t nb) {
    const int batch = (int)A.size(0), n = (int)A.size(2);
    if (n % nb != 0) throw std::runtime_error("n must be divisible by nb");
    auto L = torch::zeros_like(A);           // strict upper stays 0
    const float* pa = A.data_ptr<float>();
    float* pl = L.data_ptr<float>();
    const long long S = (long long)n * n;
    const cublasComputeType_t ct = CUBLAS_COMPUTE_32F_FAST_TF32;

    for (int k = 0; k < n; k += (int)nb) {
        int cur = (int)nb; int R = n - k;
        // copy A panel -> L panel
        {
            long long tot = (long long)R * cur;
            int th = 256; int bl = (int)((tot + th - 1) / th); if (bl > 512) bl = 512;
            dim3 grid(bl, 1, batch);
            copy_panel_kernel<<<grid, th>>>(pa, pl, n, k, cur, R);
        }
        // trailing: L[k:n,k:k+cur] = -L[k:n,0:k] @ L[k:k+cur,0:k]^T + A_panel
        if (k > 0) {
            const float* Aleft = pl + (size_t)k * n + 0;   // L[k, 0]
            float* C = pl + (size_t)k * n + k;             // L[k, k]
            gemm_rm(false, true, R, cur, k, -1.0f,
                    Aleft, n, S, Aleft, n, S, 1.0f, C, n, S, batch, ct);
        }
        // panel factor (diag + trsm)
        if (nb == 32) panel_factor_warp32<<<batch, 256>>>(pl, n, k, R);
        else if (nb == 64) panel_factor_kernel<64><<<batch, 256>>>(pl, n, k, R);
        else throw std::runtime_error("unsupported nb");
    }
    cudaError_t e = cudaGetLastError();
    if (e != cudaSuccess) throw std::runtime_error(std::string("chol_batched: ") + cudaGetErrorString(e));
    return L;
}
"""

_CPP_SRC = ("torch::Tensor chol_warp(torch::Tensor A);\n"
            "torch::Tensor chol_block(torch::Tensor A);\n"
            "torch::Tensor chol_batched(torch::Tensor A, int64_t nb);")

_MOD = load_inline(
    name="chol_ext_s2",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=["chol_warp", "chol_block", "chol_batched"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    A = data.contiguous()
    n = A.size(-1)
    if n == 32:
        return _MOD.chol_warp(A)
    if n == 64:
        return _MOD.chol_block(A)
    if n == 512:
        return _MOD.chol_batched(A, _NB)
    if n >= 1024 and A.size(0) <= 8:
        outs = [torch.linalg.cholesky_ex(A[i], check_errors=False).L
                for i in range(A.size(0))]
        return torch.stack(outs, 0)
    return torch.linalg.cholesky_ex(A, check_errors=False).L