The nvmath.device module is experimental and potentially subject to future changes.

nvmath-python Device APIs#

The device module of nvmath-python nvmath.device offers integration with NVIDIA’s high-performance computing libraries through device APIs for cuFFTDx, cuBLASDx, cuSOLVERDx, and cuRAND. Device APIs can only be called from CUDA device or kernel code, and execute on the GPU.

Users may take advantage of the device module via the approaches below:

  • numba-cuda Extensions: Users can access these device APIs via Numba by utilizing specific extensions that simplify the process of defining functions, querying device traits, and calling device functions.

  • numba-cuda-mlir Extensions: The same extensions are available with the MLIR-based Numba CUDA compiler. See Supported Compilers for the current support status.

  • Third-party JIT Compilers: The APIs are also available through low-level interfaces in other JIT compilers, allowing advanced users to work directly with the raw device code. See Using Device APIs with Third-party Compilers.

Note

The FFT, Matmul, and Solver device APIs in module nvmath.device currently support cuFFTDx 1.7.0, cuBLASDx 0.6.0, and cuSOLVERDx 0.4.0, also available as part of MathDx 26.03.0. Refer to the per-library pages below for the supported functionality. Notably, cuFFTDx C++ APIs with a workspace argument are currently not available in nvmath-python.

Supported Compilers#

The device API objects (FFT, Matmul, and Solver) can be used with both the numba-cuda and numba-cuda-mlir compilers. Switching between the compilers requires changing only the import of the cuda module:

import numpy as np
from numba_cuda_mlir import cuda
# from numba import cuda  # the only change needed to use numba-cuda

from nvmath.device import Matmul

MM = Matmul(
    size=(32, 32, 32),
    precision=np.float64,
    data_type="real",
    arrangement=("row_major", "col_major", "col_major"),
    execution="Block",
)

@cuda.jit
def kernel(alpha, a, b, beta, c):
    # stage the a, b and c matrices in shared memory
    ...
    MM.execute(alpha, smem_a, smem_b, beta, smem_c)
    # store the result back to global memory
    ...

Note

There are currently two limitations when using numba-cuda-mlir:

  • the nvmath.device.random APIs are available only with numba-cuda,

  • the advanced Matmul APIs (opaque tensors, accumulators, and pipelines) are available only with numba-cuda.

Using Device APIs with Third-party Compilers#

For JIT compilers other than the supported Numba compilers, the device functions can be obtained directly in the LTO-IR format with the compile_blas_execute, compile_fft_execute, and compile_solver_execute functions. The example below compiles the execute function of a Matmul object. It returns the LTO-IR code together with the name of the C-ABI symbol to link against:

import numpy as np

from nvmath.device import Matmul, compile_blas_execute, current_device_lto

MM = Matmul(
    size=(32, 32, 32),
    precision=np.float32,
    data_type="real",
    arrangement=("row_major", "col_major", "col_major"),
    execution="Block",
)

code, symbol = compile_blas_execute(
    MM, code_type=current_device_lto(), execute_api="static_leading_dimensions"
)

ltoir = code.data  # the LTO-IR buffer to pass to the compiler or linker

Contents