nvmath-python

nvmath-python reimagines math library APIs to cover sophisticated use cases, which are impossible with NumPy-like APIs without performance compromises. 

Host APIs provide both out-of-the-box simplicity and versatility of customization options through optional arguments for accessing all “knobs” of the underlying NVIDIA math libraries. Host APIs are divided into generic and specialized APIs. Generic APIs are intended for a consistent user experience across memory/execution spaces. They may not work with some data types, which are hardware specific and they do not necessarily leverage specific device capabilities. They are great for writing portable code. In contrast, specialized APIs have narrower applications, they may only work on a specific hardware only. They are great for fully leveraging hardware capabilities at the expense of portability.

Device APIs allow embedding nvmath-python library calls in custom kernels, written using Python compilers such as numba-cuda. You no longer need to write GEMM or FFT kernels from scratch.

Host APIs with callbacks allow embedding custom Python code into nvmath-python calls. Internal JIT machinery compiles the custom code and fuses with nvmath-python operation to achieve peak performance.

Stateful (class-form) APIs allow splitting the entire math operation into specification, planning, autotuning, and execution phases. Having an expensive specification, planning, and autotuning done once, allows their cost amortization through multiple subsequent executions.

Integration with the Python logging facility brings visibility into specification, planning, autotuning, and execution machinery details at runtime.

nvmath-python works in conjunction with popular Python packages. This includes GPU-based packages like CuPy, PyTorch, and RAPIDS and CPU-based packages like NumPy, SciPy, and scikit-learn. You can keep using familiar data structures and workflows while benefiting from accelerated math operations through nvmath-python.

nvmath-python is not a replacement for array libraries such as NumPy, CuPy, and PyTorch. It does not implement array APIs for array creation, indexing, and slicing. nvmath-python is intended to be used alongside these array libraries. All these dependencies are optional, and you’re free to choose which array library (or multiple libraries) to work with alongside nvmath-python.

nvmath-python supports CPU and GPU execution and memory spaces. It eases the transition between CPU and GPU implementations as well as allows implementing hybrid CPU-GPU workflows.

In combination with Python compilers, such as numba-cuda, you can implement GPU custom kernels with embedded nvmath-python library calls.

nvmath-python pushes performance limits to an edge by delivering performance comparable to the underlying CUDA-X native libraries, such as cuBLAS family, cuFFT family, cuDSS, and cuRAND. With stateful APIs you can amortize the costs of specification, planning, and autotuning phases through multiple executions.

For CPU execution nvmath-python leverages the NVPL library for the best performance on NVIDIA Grace™ CPU platforms. It also supports acceleration of x86 hosts by leveraging the MKL library.

In combination with Python compilers, such as numba-cuda, you can now write highly performant kernels involving GEMM, FFT, and/or RNG operations. Here are just a few examples of impossible made possible with nvmath-python:

nvmath-python allows scaling beyond a single GPU and even beyond a single node without major coding effort. Multi-GPU multi-node (MGMN) APIs allow easy transition from a single GPU implementation to MGMN and seamlessly scale to thousands of GPUs. The library also offers helper utilities for reshaping data (re-partitioning) as needed without major coding.