submission 888202

#!POPCORN leaderboard cholesky
#!POPCORN gpu B200

"""Validated cuSolverDx32 + production Triton64 hybrid experiment."""

import glob
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time
import importlib.util

import torch
import triton
import triton.language as tl
from torch.utils.cpp_extension import get_default_build_root

from task import input_t, output_t


def _unique_existing_directories(paths):
    result = []
    seen = set()
    for path in paths:
        resolved = str(Path(path).resolve())
        if resolved not in seen and Path(resolved).is_dir():
            seen.add(resolved)
            result.append(resolved)
    return result


def _find_headers(name):
    patterns = [
        f"/usr/local/cuda*/include/{name}",
        f"/usr/local/cuda*/include/cusolverdx/include/{name}",
        f"/usr/local/cuda*/include/cublasdx/include/{name}",
        f"/opt/**/{name}",
    ]
    for base in sys.path:
        if not base:
            continue
        patterns.extend(
            [
                os.path.join(base, name),
                os.path.join(base, "nvidia", "mathdx", "*", "include", name),
                os.path.join(
                    base,
                    "nvidia",
                    "mathdx",
                    "*",
                    "include",
                    "cusolverdx",
                    "include",
                    name,
                ),
                os.path.join(
                    base,
                    "nvidia",
                    "mathdx",
                    "*",
                    "include",
                    "cublasdx",
                    "include",
                    name,
                ),
            ]
        )
    matches = []
    for pattern in patterns:
        matches.extend(glob.glob(pattern, recursive=True))
    return sorted({str(Path(path).resolve()) for path in matches})


_cusolverdx_headers = _find_headers("cusolverdx.hpp")
_cublasdx_headers = _find_headers("cublasdx.hpp")
_all_headers = _cusolverdx_headers + _cublasdx_headers

_include_candidates = []
for header_string in _all_headers:
    header = Path(header_string)
    _include_candidates.append(header.parent)
    # MathDx can also expose a forwarding header in include/<library>/include.
    for parent in header.parents:
        if parent.name == "include":
            _include_candidates.append(parent)
            package_root = parent.parent
            _include_candidates.append(package_root / "external" / "cutlass" / "include")
            break
_include_paths = _unique_existing_directories(_include_candidates)

_library_patterns = [
    "/usr/local/cuda*/**/libcusolverdx.a",
    "/usr/local/cuda*/**/libcusolverdx.fatbin",
    "/opt/**/libcusolverdx.a",
    "/opt/**/libcusolverdx.fatbin",
    "/opt/**/libcublasdx.fatbin",
]
_libraries = []
for _pattern in _library_patterns:
    _libraries.extend(glob.glob(_pattern, recursive=True))
_libraries = sorted({str(Path(path).resolve()) for path in _libraries})

_nvcc = shutil.which("nvcc")
_nvcc_version = None
if _nvcc is not None:
    try:
        _nvcc_version = subprocess.run(
            [_nvcc, "--version"],
            check=True,
            capture_output=True,
            text=True,
        ).stdout.strip()
    except Exception as error:  # diagnostic only
        _nvcc_version = repr(error)

_environment = {
    "python": sys.version,
    "torch": torch.__version__,
    "torch_cuda_runtime": torch.version.cuda,
    "gpu_name": torch.cuda.get_device_name() if torch.cuda.is_available() else None,
    "gpu_capability": (
        list(torch.cuda.get_device_capability()) if torch.cuda.is_available() else None
    ),
    "torch_arch_list": torch.cuda.get_arch_list() if torch.cuda.is_available() else [],
    "nvcc": _nvcc,
    "nvcc_version": _nvcc_version,
    "cusolverdx_headers": _cusolverdx_headers,
    "cublasdx_headers": _cublasdx_headers,
    "mathdx_include_paths": _include_paths,
    "mathdx_libraries": _libraries,
}
print("CUSOLVERDX_PROBE_ENV=" + json.dumps(_environment, sort_keys=True), file=sys.stderr, flush=True)


CPP_SRC = r"""
#include <torch/extension.h>

#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#include <cstdint>
#include <limits>

extern "C" cudaError_t launch_cusolverdx_cholesky32(
    const float* input,
    float* output,
    int* info,
    int batch);
extern "C" int64_t cusolverdx_header_flags_cuda();
extern "C" int64_t cusolverdx_shared_memory_bytes_cuda();

at::Tensor cusolverdx_cholesky32(const at::Tensor& input) {
  constexpr int kN = 32;
  TORCH_CHECK(input.is_cuda(), "cuSolverDx probe expects a CUDA tensor");
  TORCH_CHECK(input.scalar_type() == at::ScalarType::Float,
              "cuSolverDx probe expects torch.float32");
  TORCH_CHECK(input.dim() == 3 && input.size(1) == kN && input.size(2) == kN,
              "cuSolverDx probe expects [batch,32,32]");
  TORCH_CHECK(input.is_contiguous(),
              "cuSolverDx probe expects contiguous input");
  TORCH_CHECK(input.size(0) <= std::numeric_limits<int>::max(),
              "cuSolverDx probe batch is too large");

  c10::cuda::CUDAGuard device_guard(input.device());
  auto output = at::empty_like(input);
  const int batch = static_cast<int>(input.size(0));
  if (batch == 0) {
    return output;
  }
  auto info = at::zeros({batch}, input.options().dtype(at::kInt));
  const cudaError_t error = launch_cusolverdx_cholesky32(
      input.data_ptr<float>(),
      output.data_ptr<float>(),
      info.data_ptr<int>(),
      batch);
  TORCH_CHECK(error == cudaSuccess,
              "cuSolverDx kernel launch failed: ", cudaGetErrorString(error));
  return output;
}

int64_t cusolverdx_header_flags() {
  return cusolverdx_header_flags_cuda();
}

int64_t cusolverdx_shared_memory_bytes() {
  return cusolverdx_shared_memory_bytes_cuda();
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def("cusolverdx_cholesky32", &cusolverdx_cholesky32);
  module.def("cusolverdx_header_flags", &cusolverdx_header_flags);
  module.def("cusolverdx_shared_memory_bytes", &cusolverdx_shared_memory_bytes);
}
"""


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

#include <cstdint>
#include <limits>

#if __has_include(<cusolverdx.hpp>)
#define POPCORN_HAS_CUSOLVERDX 1
#include <cusolverdx.hpp>
#else
#define POPCORN_HAS_CUSOLVERDX 0
#endif

#if __has_include(<cublasdx.hpp>)
#define POPCORN_HAS_CUBLASDX 1
#include <cublasdx.hpp>
#else
#define POPCORN_HAS_CUBLASDX 0
#endif

namespace {

constexpr int kN = 32;

#if POPCORN_HAS_CUSOLVERDX
using Potrf32 = decltype(
    cusolverdx::Size<kN>()
    + cusolverdx::Precision<float>()
    + cusolverdx::Type<cusolverdx::type::real>()
    + cusolverdx::Function<cusolverdx::function::potrf>()
    + cusolverdx::FillMode<cusolverdx::lower>()
    + cusolverdx::Arrangement<cusolverdx::row_major>()
    + cusolverdx::SM<1000>()
    + cusolverdx::Block());

__global__ void potrf32_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int* __restrict__ info,
    int batch) {
  const int matrix = blockIdx.x;
  if (matrix >= batch) {
    return;
  }

  extern __shared__ __align__(16) unsigned char dynamic_shared[];
  float* tile = reinterpret_cast<float*>(dynamic_shared);
  const float* matrix_input = input + matrix * kN * kN;
  float* matrix_output = output + matrix * kN * kN;

  for (int index = threadIdx.x; index < kN * kN; index += blockDim.x) {
    tile[index] = matrix_input[index];
  }
  __syncthreads();

  Potrf32().execute(tile, info + matrix);
  __syncthreads();

  for (int index = threadIdx.x; index < kN * kN; index += blockDim.x) {
    const int row = index / kN;
    const int col = index - row * kN;
    matrix_output[index] = row >= col ? tile[index] : 0.0f;
  }
}
#endif

}  // namespace

extern "C" int64_t cusolverdx_header_flags_cuda() {
  return static_cast<int64_t>(POPCORN_HAS_CUSOLVERDX)
       | (static_cast<int64_t>(POPCORN_HAS_CUBLASDX) << 1);
}

extern "C" int64_t cusolverdx_shared_memory_bytes_cuda() {
#if POPCORN_HAS_CUSOLVERDX
  return static_cast<int64_t>(Potrf32::shared_memory_size);
#else
  return -1;
#endif
}

extern "C" cudaError_t launch_cusolverdx_cholesky32(
    const float* input,
    float* output,
    int* info,
    int batch) {
#if !POPCORN_HAS_CUSOLVERDX
  return cudaErrorNotSupported;
#else
  if (batch == 0) {
    return cudaSuccess;
  }
  const dim3 threads = Potrf32::block_dim;
  const int shared_bytes = static_cast<int>(Potrf32::shared_memory_size);
  potrf32_kernel<<<batch, threads, shared_bytes>>>(
      input,
      output,
      info,
      batch);
  return cudaGetLastError();
#endif
}
"""


def _build_lto_extension():
    """Build the minimal extension with the device-link step MathDx requires."""
    static_libraries = [path for path in _libraries if path.endswith("libcusolverdx.a")]
    if not static_libraries:
        raise RuntimeError(
            "cuSolverDx header is present, but libcusolverdx.a was not found; "
            "the minimal POTRF probe requires an LTO device library"
        )

    module_name = "cholesky_cusolverdx_probe_v2"
    build_directory = Path(get_default_build_root()) / module_name
    build_directory.mkdir(parents=True, exist_ok=True)
    cpp_path = build_directory / "binding.cpp"
    cuda_path = build_directory / "kernel.cu"
    setup_path = build_directory / "setup.py"
    cpp_path.write_text(CPP_SRC)
    cuda_path.write_text(CUDA_SRC)

    library_directory = str(Path(static_libraries[0]).parent)
    setup_source = f'''\
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name={module_name!r},
    ext_modules=[
        CUDAExtension(
            name={module_name!r},
            sources={[str(cpp_path), str(cuda_path)]!r},
            include_dirs={_include_paths!r},
            library_dirs={[library_directory]!r},
            libraries=["cusolverdx"],
            dlink=True,
            dlink_libraries=["cusolverdx"],
            extra_compile_args={{
                "cxx": ["-O3", "-std=c++17"],
                "nvcc": [
                    "-O3",
                    "-std=c++17",
                    "-rdc=true",
                    "-gencode=arch=compute_100,code=lto_100",
                    "-U__CUDA_NO_HALF_OPERATORS__",
                    "-U__CUDA_NO_HALF_CONVERSIONS__",
                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
                    "-U__CUDA_NO_HALF2_OPERATORS__",
                    "-Xptxas=-v",
                ],
            }},
        )
    ],
    cmdclass={{"build_ext": BuildExtension}},
)
'''
    setup_path.write_text(setup_source)

    environment = os.environ.copy()
    environment["TORCH_CUDA_ARCH_LIST"] = "10.0"
    environment["CC"] = "gcc"
    environment["CXX"] = "g++"
    command = [
        sys.executable,
        str(setup_path),
        "build_ext",
        "--inplace",
        "--build-temp",
        str(build_directory / "objects"),
    ]
    completed = subprocess.run(
        command,
        cwd=build_directory,
        env=environment,
        capture_output=True,
        text=True,
    )
    print("CUSOLVERDX_PROBE_BUILD_COMMAND=" + json.dumps(command), file=sys.stderr)
    if completed.stdout:
        print("CUSOLVERDX_PROBE_BUILD_STDOUT_BEGIN", file=sys.stderr)
        print(completed.stdout, file=sys.stderr, end="")
        print("CUSOLVERDX_PROBE_BUILD_STDOUT_END", file=sys.stderr)
    if completed.stderr:
        print("CUSOLVERDX_PROBE_BUILD_STDERR_BEGIN", file=sys.stderr)
        print(completed.stderr, file=sys.stderr, end="")
        print("CUSOLVERDX_PROBE_BUILD_STDERR_END", file=sys.stderr)
    if completed.returncode != 0:
        raise RuntimeError(
            f"cuSolverDx LTO extension build failed with exit code "
            f"{completed.returncode}"
        )

    candidates = sorted(build_directory.glob(module_name + "*.so"))
    if len(candidates) != 1:
        raise RuntimeError(
            f"expected one built cuSolverDx extension, found {len(candidates)}"
        )
    spec = importlib.util.spec_from_file_location(module_name, candidates[0])
    if spec is None or spec.loader is None:
        raise RuntimeError("could not create import spec for cuSolverDx extension")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


_compile_start = time.perf_counter()
dx_module = _build_lto_extension()
_compile_ms = (time.perf_counter() - _compile_start) * 1.0e3
_build = {
    "compile_ms": _compile_ms,
    "header_flags": int(dx_module.cusolverdx_header_flags()),
    "cusolverdx_header_visible": bool(dx_module.cusolverdx_header_flags() & 1),
    "cublasdx_header_visible": bool(dx_module.cusolverdx_header_flags() & 2),
    "potrf32_shared_memory_bytes": int(dx_module.cusolverdx_shared_memory_bytes()),
}
print("CUSOLVERDX_PROBE_BUILD=" + json.dumps(_build, sort_keys=True), file=sys.stderr, flush=True)


@triton.jit
def _cholesky64_blocked_kernel(
    input_ptr,
    output_ptr,
    matrix_stride: tl.constexpr,
):
    matrix = tl.program_id(0)
    row_ids = tl.arange(0, 32)
    col_ids = tl.arange(0, 32)
    rows = row_ids[:, None]
    cols = col_ids[None, :]
    matrix_offset = matrix * matrix_stride
    top_offsets = matrix_offset + rows * 64 + cols

    # Factor the 32x32 leading diagonal block. Restricting each live tile to
    # 32x32 keeps register pressure much lower than a monolithic 64x64 kernel.
    top = tl.where(
        rows >= cols,
        tl.load(input_ptr + top_offsets),
        0.0,
    )

    for k in range(32):
        pivot_row = tl.sum(tl.where(rows == k, top, 0.0), axis=0)
        diagonal = tl.sum(
            tl.where(col_ids == k, pivot_row, 0.0), axis=0
        )
        diagonal -= tl.sum(
            tl.where(col_ids < k, pivot_row * pivot_row, 0.0), axis=0
        )
        diagonal = tl.sqrt(tl.maximum(diagonal, 0.0))

        column = tl.sum(tl.where(cols == k, top, 0.0), axis=1)
        products = tl.where(
            cols < k, top * pivot_row[None, :], 0.0
        )
        column = (column - tl.sum(products, axis=1)) / diagonal

        top = tl.where(
            (rows == k) & (cols == k), diagonal, top
        )
        top = tl.where(
            (rows > k) & (cols == k), column[:, None], top
        )

    # A10 = L10 @ L00.T. Solve all 32 rows of L10 together, one column at a
    # time, reusing the completed rows of L00 from the first phase.
    bottom_left_offsets = matrix_offset + (rows + 32) * 64 + cols
    bottom_left = tl.load(input_ptr + bottom_left_offsets)
    for k in range(32):
        pivot_row = tl.sum(tl.where(rows == k, top, 0.0), axis=0)
        diagonal = tl.sum(
            tl.where(col_ids == k, pivot_row, 0.0), axis=0
        )
        rhs = tl.sum(
            tl.where(cols == k, bottom_left, 0.0), axis=1
        )
        products = tl.where(
            cols < k,
            bottom_left * pivot_row[None, :],
            0.0,
        )
        solved_column = (
            rhs - tl.sum(products, axis=1)
        ) / diagonal
        bottom_left = tl.where(
            cols == k,
            solved_column[:, None],
            bottom_left,
        )

    # L00 is dead after the solve. Commit it now so the compiler can release
    # that tile before the Tensor Core update and trailing factorization.
    top_right_offsets = matrix_offset + rows * 64 + (cols + 32)
    tl.store(output_ptr + top_offsets, top)
    tl.store(output_ptr + top_right_offsets, 0.0)

    # Form the trailing Schur complement with B200 Tensor Cores. tf32x3 uses
    # three TF32 products to recover near-FP32 accuracy before FP32 Cholesky.
    bottom_right_offsets = (
        matrix_offset + (rows + 32) * 64 + (cols + 32)
    )
    schur = tl.load(input_ptr + bottom_right_offsets)
    schur -= tl.dot(
        bottom_left,
        tl.trans(bottom_left),
        input_precision="tf32x3",
        out_dtype=tl.float32,
    )
    # The dot is L10's final consumer; write it before factoring the Schur
    # complement to avoid keeping another 32x32 tile live.
    tl.store(output_ptr + bottom_left_offsets, bottom_left)
    trailing = tl.where(rows >= cols, schur, 0.0)

    for k in range(32):
        pivot_row = tl.sum(
            tl.where(rows == k, trailing, 0.0), axis=0
        )
        diagonal = tl.sum(
            tl.where(col_ids == k, pivot_row, 0.0), axis=0
        )
        diagonal -= tl.sum(
            tl.where(col_ids < k, pivot_row * pivot_row, 0.0), axis=0
        )
        diagonal = tl.sqrt(tl.maximum(diagonal, 0.0))

        column = tl.sum(
            tl.where(cols == k, trailing, 0.0), axis=1
        )
        products = tl.where(
            cols < k, trailing * pivot_row[None, :], 0.0
        )
        column = (column - tl.sum(products, axis=1)) / diagonal

        trailing = tl.where(
            (rows == k) & (cols == k), diagonal, trailing
        )
        trailing = tl.where(
            (rows > k) & (cols == k), column[:, None], trailing
        )

    tl.store(output_ptr + bottom_right_offsets, trailing)


def custom_kernel(data: input_t) -> output_t:
    batch, n, _ = data.shape
    if n == 32:
        return dx_module.cusolverdx_cholesky32(data)
    if n == 64:
        output = torch.empty_like(data)
        _cholesky64_blocked_kernel[(batch,)](
            data,
            output,
            n * n,
            num_warps=2,
        )
        return output
    if batch == 2 and n == 4096:
        output = torch.empty_like(data)
        info = torch.empty((2,), dtype=torch.int32, device=data.device)
        torch.linalg.cholesky_ex(
            data[0],
            check_errors=False,
            out=(output[0], info[0]),
        )
        torch.linalg.cholesky_ex(
            data[1],
            check_errors=False,
            out=(output[1], info[1]),
        )
        return output
    return torch.linalg.cholesky_ex(data, check_errors=False).L