submission 881304

#!POPCORN leaderboard cholesky
#!POPCORN gpu B200

import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t


CPP_SRC = r"""
torch::Tensor cholesky_cuda(torch::Tensor input);
"""

CUDA_SRC = "#include <torch/extension.h>\n#include <ATen/Functions.h>\n#include <cublas_v2.h>\n#include <cusolverDn.h>\n\n#include <algorithm>\n#include <cfloat>\n#include <tuple>\n\n// for small matrices, each block for 1 matrix in batch\n__global__ void cholesky_shared(const float* __restrict__ input, float* __restrict__ output, int n) {\n\n  // full matrix stored in shared memory\n  extern __shared__ float tile[];\n  const int pitch = n + 1; // better bank conflicts, extra column\n  const int count = n * n;\n  const long long offset = static_cast<long long>(blockIdx.x) * count;\n  const float* a = input + offset;\n  float* l = output + offset;\n\n  // load lower half\n  for (int index = threadIdx.x; index < count; index += blockDim.x) {\n    const int row = index / n;\n    const int col = index - row * n;\n    if(col <= row){\n      tile[row * pitch + col] = a[index];\n    } else {\n      tile[row * pitch + col] = 0.0f;\n    }\n  }\n  __syncthreads();\n\n  // each col\n  for (int k = 0; k < n; ++k) {\n    \n    // diagonal: \n    if (threadIdx.x == 0) {\n      float pivot = tile[k * pitch + k];\n      for (int j = 0; j < k; ++j) {\n        const float x = tile[k * pitch + j];\n        pivot = fmaf(-x, x, pivot); // this area we can swap the adds for a reduction like kernel ()?\n      }\n      tile[k * pitch + k] = sqrtf(fmaxf(pivot, FLT_MIN));\n    }\n    __syncthreads();\n\n    const float diagonal = tile[k * pitch + k];\n\n    // below diagonal\n    for (int row = k + 1 + threadIdx.x; row < n; row += blockDim.x) {\n      float value = tile[row * pitch + k];\n      for (int j = 0; j < k; ++j) {\n        value = fmaf(-tile[row * pitch + j], tile[k * pitch + j], value);\n      }\n      tile[row * pitch + k] = value / diagonal;\n    }\n    __syncthreads();\n  }\n\n  // write to output\n  for (int index = threadIdx.x; index < count; index += blockDim.x) {\n    const int row = index / n;\n    const int col = index - row * n;\n    l[index] = tile[row * pitch + col];\n  }\n}\n\n__global__ void zero_upper(float* output, long long count, int n) {\n  for (long long index = static_cast<long long>(blockIdx.x) * blockDim.x + threadIdx.x;\n       index < count;\n       index += static_cast<long long>(blockDim.x) * gridDim.x) {\n    const int row = static_cast<int>((index / n) % n);\n    const int col = static_cast<int>(index % n);\n    if (col > row) output[index] = 0.0f;\n  }\n}\n\n// cuBLAS handle for triangular panel solve and trailing symmetric update \ncublasHandle_t blas_handle() {\n  static cublasHandle_t handle = [] {\n    cublasHandle_t value;\n    cublasCreate(&value);\n    cublasSetMathMode(value, CUBLAS_DEFAULT_MATH);\n    return value;\n  }();\n  return handle;\n}\n\n// cuSOLVER handle for dn spotrf\ncusolverDnHandle_t solver_handle() {\n  static cusolverDnHandle_t handle = [] {\n    cusolverDnHandle_t value;\n    cusolverDnCreate(&value);\n    return value;\n  }();\n  return handle;\n}\n\n// blocked cholesky (right looking)\ntorch::Tensor cholesky_blocked(torch::Tensor input) {\n  constexpr int tile_size = 512;\n  const int n = static_cast<int>(input.size(1));\n  auto output = input.clone();\n  float* base = output.data_ptr<float>();\n  cublasHandle_t blas = blas_handle();\n  cusolverDnHandle_t solver = solver_handle();\n\n  int workspace_size = 0;\n  cusolverDnSpotrf_bufferSize(solver, CUBLAS_FILL_MODE_UPPER, std::min(tile_size, n), base, n, &workspace_size);\n  auto workspace = torch::empty({workspace_size}, input.options());\n  auto info = torch::empty({1}, input.options().dtype(torch::kInt32));\n  const float one = 1.0f;\n  const float minus_one = -1.0f;\n\n  // iterate over tiles, each iter does 1 diag and everything below\n  for (int k = 0; k < n; k += tile_size) {\n    const int block = std::min(tile_size, n - k);\n    const int trailing = n - k - block;\n    float* diagonal = base + static_cast<long long>(k) * n + k;\n\n    // factor diagonal \n    cusolverDnSpotrf(solver, CUBLAS_FILL_MODE_UPPER, block, diagonal, n, workspace.data_ptr<float>(), workspace_size, info.data_ptr<int>());\n\n    if (trailing == 0) continue;\n\n    float* panel = base + static_cast<long long>(k + block) * n + k;\n    float* remainder = base + static_cast<long long>(k + block) * n + k + block;\n\n    // panel below diagonal\n    cublasStrsm(blas, CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, block, trailing, &one, diagonal, n, panel, n);\n\n    // updates remaining matrix\n    cublasSsyrk(blas, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_T, trailing, block, &minus_one, panel, n, &one, remainder, n);\n  }\n\n  const long long count = output.numel();\n  constexpr int threads = 256;\n  const int blocks = static_cast<int>(std::min<long long>((count + threads - 1) / threads, 4096));\n  zero_upper<<<blocks, threads>>>(base, count, n);\n  return output;\n}\n\ntorch::Tensor cholesky_cuda(torch::Tensor input) {\n  const int batch = static_cast<int>(input.size(0));\n  const int n = static_cast<int>(input.size(1));\n\n  if (n <= 128) {\n    static const cudaError_t configured = cudaFuncSetAttribute(cholesky_shared, cudaFuncAttributeMaxDynamicSharedMemorySize, 128 * 129 * static_cast<int>(sizeof(float)));\n    (void)configured;\n    auto output = torch::empty_like(input);\n    const int threads = n <= 32 ? 32 : (n <= 64 ? 64 : 128);\n    const size_t shared = static_cast<size_t>(n) * (n + 1) * sizeof(float);\n    cholesky_shared<<<batch, threads, shared>>>(input.data_ptr<float>(), output.data_ptr<float>(), n);\n    return output;\n  }\n\n  if (batch == 1 && n >= 2048) return cholesky_blocked(input);\n\n  return std::get<0>(at::linalg_cholesky_ex(input, false, false));\n}\n"

module = load_inline(
    name="cholesky_cuda_extension",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["cholesky_cuda"],
    extra_cuda_cflags=["-O3", "--std=c++17"],
    extra_ldflags=["-lcublas", "-lcusolver"],
    with_cuda=True,
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    return module.cholesky_cuda(data)