cuSolverDx Python Bindings

Overview

cuSolverDx offers C++ APIs that are callable from CUDA C++ kernels, but its functionalities can also be easily accessed from Python using NVIDIA Warp or nvmath-python.

Note

cuSolverDx Python bindings match C++ code in performance and offer extensive and easy autotuning opportunities by bypassing the CUDA C++ compilation limitation.

NVIDIA Warp Integration

NVIDIA Warp is a Python library that allows developers to write high-performance simulation and graphics code that runs efficiently on both CPUs and NVIDIA GPUs. It uses just-in-time (JIT) compilation to turn Python functions into fast, parallel kernels, making it ideal for tasks like physics simulation, robotics, and geometry processing. Warp also supports differentiable programming, allowing integration with machine learning frameworks for gradient-based optimization - all while keeping the simplicity of Python.

Warp uses cuSolverDx for its linear solver operations in Tile mode.

Example Implementation

This is what a simple Cholesky factor and solve kernel using Warp looks like:

@wp.kernel
def cholesky(
    A: wp.array2d(dtype=wp_type),
    L: wp.array2d(dtype=wp_type),
    X: wp.array1d(dtype=wp_type),
    Y: wp.array1d(dtype=wp_type),
):

    i, j, _ = wp.tid()

    a = wp.tile_load(A, shape=(TILE, TILE))
    l = wp.tile_cholesky(a)
    wp.tile_store(L, l)

    x = wp.tile_load(X, shape=TILE)
    y = wp.tile_cholesky_solve(l, x)
    wp.tile_store(Y, y)

Additional Resources

The Warp GitHub repository can be accessed at NVIDIA/warp and offers multiple examples, including a simple example with warp Cholesky Tile APIs using cuSolverDx’s Cholesky factorization and triangular solver, or a blocked Cholesky example.

Warp provides autotuning out of the box through its Tile model of programming, where the user describes the problem on a high level, and then it is autotuned and mapped onto the hardware by Warp.

nvmath-python Integration

nvmath-python is a Python library that provides high-performance, pythonic access to NVIDIA’s CUDA-X math libraries, enabling accelerated mathematical operations like linear algebra and fast Fourier transforms on both CPUs and NVIDIA GPUs. It integrates seamlessly with popular Python libraries such as CuPy, PyTorch, and NumPy, allowing users to leverage NVIDIA hardware acceleration within familiar workflows without needing C or C++ bindings. With both stateless and stateful APIs, nvmath-python delivers near-native performance for computational tasks in deep learning, data processing, and scientific computing, while supporting advanced features like device kernel fusion and customizable callbacks.

Example Implementation

This is what a simple Cholesky factorization kernel using nvmath-python looks like:

solver = CholeskySolver(
    size=(32, 32),
    precision=np.float64,
    data_type="complex",
    execution="Block",
    leading_dimensions=(33, 33),
    block_dim=(256, 1, 1),
    fill_mode="upper",
)

n_a = solver.a_size()

@cuda.jit
def kernel(a, info):
    smem_a = cuda.shared.array(n_a, dtype=solver.value_type)

    load_to_shared_strided(a, smem_a, solver.a_shape, solver.a_strides())
    cuda.syncthreads()

    solver.factorize(smem_a, info)
    cuda.syncthreads()

    store_from_shared_strided(smem_a, a, solver.a_shape, solver.a_strides())

Additional Resources

The nvmath-python GitHub repository can be accessed at NVIDIA/nvmath-python and offers multiple examples, including a simple example with nvmath-python Cholesky device APIs using cuSolverDx’s Cholesky factorization, or an nvmath-python blocked Cholesky example.