Effortless Interoperability#

All host APIs support input arrays/tensors from NumPy, CuPy, and PyTorch while returningoutput operands using the same package, thus offering effortless interoperability with theseframeworks. One example for the interoperability is shown below:

importnumpyasnpimportnvmath# Create a numpy.ndarray as inputa=np.random.random(128)+1.j*np.random.random(128)# Call nvmath-python Pythonic APIsb=nvmath.fft.fft(a)# Verify that output is also a numpy.ndarrayassertisinstance(b,np.ndarray)

Stateless and Stateful APIs#

The host APIs within nvmath-python can be generally categorized into two types: statelessfunction-form APIs and stateful class-form APIs.

The function-form APIs, such asnvmath.fft.fft() andnvmath.linalg.advanced.matmul(), are designed to deliver quick, end-to-end resultswith a single function call. These APIs are ideal for instances where a user needs toperform a single computation without the need for intermediate steps, customization ofalgorithm selection, or cost amortization of preparatory steps. Conversely, the statefulclass-form APIs, likenvmath.fft.FFT andnvmath.linalg.advanced.Matmul,offer a more comprehensive and flexible approach. They not only encompass the functionalityfound in their function-form counterparts but also allow for amortization of one-time costs,potentially enhancing performance significantly.

The design pattern for all stateful APIs in nvmath-python consists of several key phases:

• Problem Specification: This initial phase involves defining the operation and settingoptions that affect its execution. It’s designed to be as lightweight as possible,ensuring the problem is well-defined and supported by the current implementation.

• Preparation: Using FFT as an example, this phase includes a planning step to selectthe optimal algorithm for the defined FFT operation. An optional autotuning operation,when available, also falls within the preparation phase. The preparation phase isgenerally the most resource-intensive and may incorporate user-specified planning andautotuning options.

• Execution: This phase allows for repeated execution, where the operand can be eithermodified in-place or explicitly reset using thereset_operand/reset_operandsmethod. The costs associated with the first two phases are therefore amortized overthese multiple executions.

• Resource Release: Users are advised to use stateful objects from within a contextusing thewith statement, whichautomatically handles the release of internal resources upon exit. If the object isnot used as a context manager usingwith, it is necessary to explicitly call thefree method to ensure all resources are properly released.

Generic and Specialized APIs#

Another way of categorizing the host APIs within nvmath-python is by splitting them intogeneric andspecialized APIs, based on their flexibility and the scope of theirfunctionality:

• Generic APIs are designed to accommodate a broad range of operands and customizationwith these APIs is confined to options that are universally applicable across allsupported operand types. For instance, the generic matrix multiplication API can handlestructured matrices (such as triangular and banded, in full or packed form) in addition todense full matrices, but the available options are limited to those applicable to allthese matrix types.

• Specialized APIs, on the other hand, are tailored for specific types of operands,allowing for full customization that is available to this kind. A prime example is thespecialized matrix multiplication API for dense matrices, which provides numerous optionsspecifically suited to dense matrices.

It should be noted that the notion of generic and specialized APIs is orthogonal to thenotion of stateful versus stateless APIs. Currently, nvmath-python offers the specializedinterface for dense matrix multiplication, instateful andstatelessforms.

Full Logging Support#

nvmath-python provides integration with the Python standard library logger from theloggingmodule to offer full logging of thecomputational details at various levels, for example debug, information, warning and error.An example illustrating the use of the global Python logger is shown below:

importlogging# Turn on logging with level set to "debug" and use a custom format for the loglogging.basicConfig(level=logging.DEBUG,format='%(asctime)s%(levelname)-8s%(message)s',datefmt='%m-%d %H:%M:%S')# Call nvmath-python Pythonic APIsout=nvmath.linalg.advanced.matmul(...)

Alternatively, for APIs that contain theoptions argument, users can set a custom loggerby directly passing it inside a dictionary or as part of the correspondingOptionsobject, for examplenvmath.fft.FFTOptions.logger fornvmath.fft.fft() andnvmath.fft.FFT. An example based on FFT is shown below:

importlogging# Create a custom loggerlogger=logging.getLogger('userlogger')...# Call nvmath-python Pythonic APIsout=nvmath.fft.fft(...,options={'logger':logger})

For the complete examples, refer toglobal logging exampleandcustom user logging example.

Call Blocking Behavior#

By default, calls to all Pythonic host APIs that require GPU execution arenot blocking ifthe input operands reside on the device. This means that functions likenvmath.linalg.advanced.matmul(),nvmath.fft.FFT.execute(), andnvmath.linalg.advanced.Matmul.execute() will return immediately after the operation islaunched on the GPU without waiting for it to complete. Users are therefore responsible forproperly synchronizing the stream when needed. The default behavior can be modified bysetting theblocking attribute (default'auto') of the relevantOptions objecttoTrue. For example, users may setnvmath.fft.FFTOptions.blocking toTrueand pass this options object to the corresponding FFT API calls. If the input operands areon the host, the Pythonic API calls will always block since the computation yields an outputoperand that will also reside on the host. Meanwhile, APIs that execute on the host (such asnvmath.fft.FFT.create_key()) always block.

Stream Semantics#

The stream semantics depend on whether the behavior of the execution APIs is chosen to beblocking or non-blocking (seeCall Blocking Behavior).

For blocking behavior, stream ordering is automatically handled by the nvmath-pythonhigh-level APIs foroperations that are performed within the package. A stream can beprovided for two reasons:

1. When the computation that prepares the input operands is not already complete by the timethe execution APIs are called. This is a correctness requirement for user-provided data.

2. To enable parallel computations across multiple streams if the device has sufficientresources and the current stream (which is the default) has concomitant operations. Thiscan be done for performance reasons.

For non-blocking behavior, it is the user’s responsibility to ensure correct stream orderingbetween the execution API calls.

The execution APIs are always launched on the provided stream.

For examples on stream ordering, refer toFFT with multiple streams.

Memory Management#

By default, the host APIs use the memory pool from the package that their operands belongto. This ensures that there is no contention for memory or spurious out-of-memory errors.However, the user also has the ability to provide their own memory allocator if they chooseto do so. In our Pythonic APIs, we support anEMM-like interface as proposed andsupported by Numba for users to set their Python mempool. Taking FFT as an example, userscan set the optionnvmath.fft.FFTOptions.allocator to a Python object complying withthenvmath.BaseCUDAMemoryManager protocol, and pass the options to the high-levelAPIs likenvmath.fft.fft() ornvmath.fft.FFT. Temporary memory allocationswill then be done through this interface. Internally, we use the same interface to use CuPyor PyTorch’s mempool depending on the operands.