submission 893127

#!POPCORN leaderboard cholesky
#!POPCORN gpu B200
# AUTO-GENERATED by scripts/build_submission.py — edit the sources in the problem's source directory, not here.

# Host / Python side of the submission. Edit this file for anything that is
# not raw kernel code. The C++/CUDA lives in the sibling .cu / .cpp files and
# is pulled in below.
#
# cuSOLVERDx is not header-only: its routines ship as LTO-IR inside
# libcusolverdx.fatbin and must be device-LTO-linked, which load_inline()'s JIT
# build can't do. So we drive nvcc ourselves (compile the device kernel with
# -dc -dlto, then `nvcc -dlink -dlto` against the fatbin) and hand the resulting
# object files to load_inline for its ordinary host link. nvcc is available in
# the environment (load_inline itself invokes it).
#
# Run this file directly for local development: the block between the POPCORN
# SOURCES markers loads the kernels from disk. `scripts/build_submission.py`
# replaces that block with embedded string literals to produce the single-file,
# submittable ../submission.py.

import os
import subprocess
import tempfile

import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t


# >>> POPCORN SOURCES
# Kernel sources embedded by scripts/build_submission.py. Edit them in the source directory, not here.
CPP_SRC = r"""// C++ binding declaration. This signature is exposed to Python via
// load_inline()'s `functions=[...]` list and must match the definition in
// cholesky_kernel.cu exactly.

#include <torch/torch.h>

torch::Tensor cholesky_forward(torch::Tensor A);
int64_t cholesky_required_smem();
"""
CUDA_SRC = r"""// Torch-side wrapper for the cuSOLVERDx Cholesky kernel.
//
// This IS compiled by load_inline. It contains no device code -- it just
// unpacks the tensor and calls the extern "C" launcher that lives in
// cholesky_device.o (compiled + device-linked separately by submission.py and
// handed to load_inline via extra_ldflags). Keeping the torch-dependent code
// here means load_inline controls its ABI/flags, while the cuSOLVERDx device
// code was device-LTO-linked by our own nvcc.

#include <torch/torch.h>

extern "C" void cusolverdx_potrf_launch(void *A, int *info, int batch);
extern "C" int cusolverdx_potrf_n();
extern "C" unsigned cusolverdx_potrf_smem();

// Shared-memory bytes the compiled kernel needs; Python compares this against
// B200's fixed per-block limit to decide whether to fall back to torch.
int64_t cholesky_required_smem() {
  return static_cast<int64_t>(cusolverdx_potrf_smem());
}

torch::Tensor cholesky_forward(torch::Tensor A) {
  TORCH_CHECK(A.is_cuda(), "input must be a CUDA tensor");
  TORCH_CHECK(A.dim() >= 2, "input must be at least 2-D ([..., N, N])");
  TORCH_CHECK(A.scalar_type() == torch::kFloat32,
              "only float32 supported, got ", A.scalar_type());
  const int64_t n = A.size(-1);
  TORCH_CHECK(A.size(-2) == n, "trailing dims must be square, got ", A.size(-2),
              "x", n);

  // The device object was compiled for a fixed size; submission.py picks the
  // matching build, but assert it here so a mismatch fails loudly.
  TORCH_CHECK(n == cusolverdx_potrf_n(),
              "kernel built for N=", cusolverdx_potrf_n(), " but got N=", n);

  auto out = A.contiguous().clone();
  const int64_t batch = out.numel() / (n * n);
  auto info = torch::zeros({batch}, A.options().dtype(torch::kInt32));

  cusolverdx_potrf_launch(out.data_ptr(), info.data_ptr<int32_t>(),
                          static_cast<int>(batch));
  return out;
}
"""
DEVICE_SRC = r"""// cuSOLVERDx device kernel for Cholesky (potrf).
//
// This TU is NOT compiled by load_inline. submission.py compiles it with our
// own nvcc using `-dc -dlto` and then `nvcc -dlink -dlto <this>.o
// libcusolverdx.fatbin` -- the device-LTO link that pulls cuSOLVERDx's LTO-IR
// out of the fatbin. load_inline's JIT build never emits that device-link step,
// so we do it here and hand the finished object files to load_inline for its
// ordinary host link.
//
// The interface to the torch side is a plain C ABI (extern "C", raw pointers)
// so there are no C++ ABI concerns across the two compilers. Inputs are
// float32; N_SIZE / SOLVER_SM are injected as -D defines at build time.
//
// NOTE: the leaderbot uploader rejects any submission text containing the
// substring "s-t-r-e-a-m", so this file avoids that word and launches on the
// default queue (the harness times there anyway).

#include <cuda_runtime.h>
#include <cusolverdx.hpp>

#ifndef N_SIZE
#define N_SIZE 32
#endif
#ifndef SOLVER_SM
#define SOLVER_SM 800
#endif

using namespace cusolverdx;

using Solver =
    decltype(Size<N_SIZE, N_SIZE>() + Precision<float>() + Type<type::real>() +
             Function<potrf>() + FillMode<fill_mode::lower>() +
             SM<SOLVER_SM>() + Block());

using DataT = typename Solver::a_data_type;

// One symmetric-positive-definite matrix per block, factored in shared memory.
__global__
__launch_bounds__(Solver::max_threads_per_block) void potrf_kernel(DataT *A,
                                                                   int *info) {
  extern __shared__ cusolverdx::byte smem[];
  DataT *As = reinterpret_cast<DataT *>(smem);

  constexpr int n = Solver::m_size;
  constexpr int lda = Solver::lda; // library-chosen shared-memory ld
  const int tid = threadIdx.x;
  const int nthread = blockDim.x;
  DataT *Ag = A + static_cast<size_t>(blockIdx.x) * n * n;

  // Row-major gmem -> column-major smem (symmetric input, so orientation is
  // moot).
  for (int k = tid; k < n * n; k += nthread) {
    int i = k / n, j = k % n;
    As[i + j * lda] = Ag[i * n + j];
  }
  __syncthreads();

  Solver().execute(As, info + blockIdx.x);
  __syncthreads();

  // Write the lower factor back row-major; zero the strict upper triangle.
  for (int k = tid; k < n * n; k += nthread) {
    int i = k / n, j = k % n;
    Ag[i * n + j] = (i >= j) ? As[i + j * lda] : DataT(0);
  }
}

extern "C" {

// So the host wrapper can assert it matches this build.
int cusolverdx_potrf_n() { return N_SIZE; }

// Shared-memory bytes this kernel needs (the whole matrix lives in smem).
unsigned cusolverdx_potrf_smem() { return Solver::shared_memory_size; }

// Launch `batch` independent factorizations on the default queue. `A` is
// [batch, N, N] row-major (contiguous); `info` is int[batch] (0 => success).
void cusolverdx_potrf_launch(void *A, int *info, int batch) {
  static bool configured = false;
  if (!configured) {
    if (Solver::shared_memory_size >
        48 * 1024) { // opt in above the 48 KB static cap
      cudaFuncSetAttribute(potrf_kernel,
                           cudaFuncAttributeMaxDynamicSharedMemorySize,
                           Solver::shared_memory_size);
    }
    configured = true;
  }
  potrf_kernel<<<batch, Solver::block_dim, Solver::shared_memory_size>>>(
      static_cast<DataT *>(A), info);
}

} // extern "C"
"""
# <<< POPCORN SOURCES

# B200 (sm_100) max opt-in dynamic shared memory per block: 227 KiB. cuSOLVERDx
# block potrf stages the whole matrix in shared memory, so a build whose
# requirement exceeds this can't launch -- we fall back to torch for those.
B200_MAX_SMEM = 227 * 1024

MATHDX_HOME = os.environ.get("MATHDX_HOME", "/opt/mathdx")
CUSOLVERDX_FATBIN = os.environ.get(
    "CUSOLVERDX_FATBIN", os.path.join(MATHDX_HOME, "lib", "libcusolverdx.fatbin")
)
NVCC = os.environ.get("CUDA_NVCC", "nvcc")


def _run(cmd):
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(
            "command failed:\n  "
            + " ".join(cmd)
            + "\n--- stdout ---\n"
            + proc.stdout
            + "\n--- stderr ---\n"
            + proc.stderr
        )


def _device_link(n, workdir):
    """Compile the cuSOLVERDx kernel to LTO-IR and device-link it against
    libcusolverdx.fatbin -- the step load_inline can't do -- returning the two
    host object files for load_inline's final link."""
    major, minor = torch.cuda.get_device_capability()
    cc = major * 10 + minor
    arch = os.environ.get(
        "CUSOLVERDX_ARCH", f"sm_{cc}"
    )  # e.g. Blackwell may want sm_100a

    src = os.path.join(workdir, "cholesky_device.cu")
    with open(src, "w") as f:
        f.write(DEVICE_SRC)
    obj = os.path.join(workdir, "cholesky_device.o")
    dlink = os.path.join(workdir, "cholesky_device_dlink.o")

    # Object files land in a shared library, so host parts must be -fPIC.
    common = [
        "-std=c++17",
        f"-arch={arch}",
        "-dlto",
        "--expt-relaxed-constexpr",
        "--extended-lambda",
        "-Xcompiler",
        "-fPIC",
        f"-I{MATHDX_HOME}/include",
        f"-I{MATHDX_HOME}/external/cutlass/include",
    ]
    # 1) compile to a relocatable LTO-IR object
    _run(
        [
            NVCC,
            "-dc",
            *common,
            f"-DN_SIZE={n}",
            f"-DSOLVER_SM={cc * 10}",
            src,
            "-o",
            obj,
        ]
    )
    # 2) device-link with the cuSOLVERDx fatbin -> object carrying the cubin
    _run(
        [
            NVCC,
            "-dlink",
            f"-arch={arch}",
            "-dlto",
            "-Xcompiler",
            "-fPIC",
            obj,
            CUSOLVERDX_FATBIN,
            "-o",
            dlink,
        ]
    )
    return [obj, dlink]


# cuSOLVERDx fixes the matrix size at compile time (inputs are float32), so we
# build (and cache) one module per N the first time we see it.
_modules = {}


def _get_module(n):
    if n not in _modules:
        tag = f"{n}_f32"
        workdir = os.path.join(tempfile.gettempdir(), f"cholesky_dx_{tag}")
        os.makedirs(workdir, exist_ok=True)
        objs = _device_link(n, workdir)
        _modules[n] = load_inline(
            name=f"cholesky_dx_{tag}",
            cpp_sources=[CPP_SRC],
            cuda_sources=[CUDA_SRC],
            functions=["cholesky_forward", "cholesky_required_smem"],
            verbose=True,
            # Hand the pre-device-linked objects to load_inline's host link.
            extra_ldflags=[*objs, "-lcudadevrt"],
        )
    return _modules[n]


known_shape_n = (32, 64, 128)
for n in known_shape_n:
    _get_module(n)


def _torch_cholesky(data):
    return torch.linalg.cholesky_ex(data, check_errors=False).L


def custom_kernel(data: input_t) -> output_t:
    n = int(data.shape[-1])
    # The matrix alone (float32) already overflows shared memory -- don't even
    # build the cuSOLVERDx kernel, just use torch.
    if n > 128:
        return _torch_cholesky(data)
    module = _get_module(n)
    return module.cholesky_forward(data)