Matmul#

Create a stateful object encapsulating the specified matrix multiplication computation \(\alpha a @ b + \beta c\) and the required resources to perform the operation. A stateful object can be used to amortize the cost of preparation (planning in the case of matrix multiplication) across multiple executions (also see the Stateful APIs section).

The function-form API matmul() is a convenient alternative to using stateful objects for single use (the user needs to perform just one matrix multiplication, for example), in which case there is no possibility of amortizing preparatory costs. The function-form APIs are just convenience wrappers around the stateful object APIs.

Using the stateful object typically involves the following steps:

Problem Specification: Initialize the object with a defined operation and options.
Preparation: Use plan() to determine the best algorithmic implementation for this specific matrix multiplication operation.
Execution: Perform the matrix multiplication computation with execute().

Detailed information on what’s happening in the various phases described above can be obtained by passing in a logging.Logger object to MatmulOptions or by setting the appropriate options in the root logger object, which is used by default:

>>> import logging
>>> logging.basicConfig(
...     level=logging.INFO,
...     format="%(asctime)s %(levelname)-8s %(message)s",
...     datefmt="%m-%d %H:%M:%S",
... )

A user can select the desired logging level and, in general, take advantage of all of the functionality offered by the Python logging module.

Parameters:

a – A tensor or a sequence of tensors representing the first operand to the matrix multiplication (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
b – A tensor or a sequence of tensors representing the second operand to the matrix multiplication (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
c – (Optional) A tensor or a sequence of tensors representing the operand to add to the matrix multiplication result (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
alpha – The scale factor or a sequence of scale factors for the matrix multiplication term as a real or complex number. The default is \(1.0\).
beta – The scale factor or a sequence of scale factors for the matrix addition term as a real or complex number. A value for beta must be provided if operand c is specified.
qualifiers – If desired, specify the matrix qualifiers as a numpy.ndarray of matrix_qualifiers_dtype objects of length <= 3 corresponding to the operands a, b, and c. By default, GeneralMatrixQualifier is assumed for each tensor. See Matrix and Tensor Qualifiers for the motivation behind qualifiers.
options – Specify options for the matrix multiplication as a MatmulOptions object. Alternatively, a dict containing the parameters for the MatmulOptions constructor can also be provided. If not specified, the value will be set to the default-constructed MatmulOptions object.
execution – Specify execution space options for the Matmul as a ExecutionCUDA or ExecutionCPU object. If not specified, the execution space will be selected to match operand’s storage (in GPU or host memory), and the corresponding ExecutionCUDA or ExecutionCPU object will be default-constructed.
stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

Semantics:

Inputs may be batched either implicitly and/or explicitly.

An implicitly-batched input is a higher-dimensional \(N \geq 3D\) tensor.
An explicitly-batched input is a Python sequence of tensors.

There may be only one explicitly-batched dimension, but there may be any number of implicitly-batched dimensions. Explicitly-batched inputs are first broadcast to equal length along the explicitly-batched dimension, then each explicit batch is treated as implicitly-batched inputs.

The semantics of the matrix multiplication follows numpy.matmul semantics, with some restrictions.

Broadcasting along implicitly-batched dimensions of c is not supported in this API, but may be supported in the future. See the advanced API (nvmath.linalg.advanced.matmul()) for an API that supports broadcasting c.

In addition, the semantics for the fused matrix addition are described below:

If arguments a and b are matrices, they are multiplied according to the rules of matrix multiplication.
If argument a is 1-D, it is promoted to a matrix by prefixing 1 to its dimensions. After matrix multiplication, the prefixed 1 is removed from the result’s dimensions.
If argument b is 1-D, it is promoted to a matrix by appending 1 to its dimensions. After matrix multiplication, the appended 1 is removed from the result’s dimensions.
The operand for the matrix addition c must be the expected shape of the result of the matrix multiplication.

Thread Safety:

A Matmul instance should only be used by its creating thread; it is not safe to use concurrently from multiple threads. The free function matmul() creates a fresh instance per call and is safe to call concurrently from multiple threads.

Examples

>>> import numpy as np
>>> import nvmath

Create two 2-D float64 ndarrays on the CPU:

>>> M, N, K = 1024, 1024, 1024
>>> a = np.random.rand(M, K)
>>> b = np.random.rand(K, N)

We will define a matrix multiplication operation using the generic matrix multiplication interface.

Create a Matmul object encapsulating the problem specification above:

>>> mm = nvmath.linalg.Matmul(a, b)

Options can be provided above to control the behavior of the operation using the options argument (see MatmulOptions).

Next, plan the operation. The operands’ layouts, qualifiers, and dtypes will be considered to select an appropriate matrix multiplication:

>>> mm.plan()

Now execute the matrix multiplication, and obtain the result r1 as a NumPy ndarray.

>>> r1 = mm.execute()

Note that all Matmul methods execute on the current stream by default. Alternatively, the stream argument can be used to run a method on a specified stream.

Let’s now look at the same problem with CuPy ndarrays on the GPU.

Create a 2-D CuPy ndarray on the GPU:

>>> import cupy as cp
>>> a = cp.random.rand(M, K)
>>> b = cp.random.rand(K, N)

Create an Matmul object encapsulating the problem specification described earlier and use it as a context manager.

>>> with nvmath.linalg.Matmul(a, b) as mm:
...     # Plan the operation.
...     mm.plan()
...
...     # Execute the operation to get the first result.
...     r1 = mm.execute()
...
...     # Update operands 'a' and 'b' in-place (see reset_operands() for an
...     # alternative).
...     a[:] = cp.random.rand(M, K)
...     b[:] = cp.random.rand(K, N)
...
...     # Execute the operation to get the new result.
...     r2 = mm.execute()

All the resources used by the object are released at the end of the block.

Further examples can be found in the examples/linalg/generic/matmul directory.

Attributes

options#

valid_state#

Methods

execute( *, stream: AnyStream | int | None = None, ) → AnyTensor | Sequence[AnyTensor][source]#

Execute a prepared (planned) matrix multiplication.

Parameters:: stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.
Returns:: The result of the specified matrix multiplication, which remains on the same device and belongs to the same package as the input operands.

execute_unchecked( *, stream: AnyStream | int | None = None, ) → AnyTensor | Sequence[AnyTensor][source]#

This method is experimental and potentially subject to future changes.

Added in version 1.0.

This method is a performance-optimized alternative to execute() that eliminates validation and logging overhead, making it ideal for performance-critical loops where operand compatibility is guaranteed by the caller.

Parameters:: stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.
Returns:: The result of the specified matrix multiplication, which remains on the same device and belongs to the same package as the input operands.

See also

execute(): Safe, validated method for executing the matrix multiplication.

free()[source]#

Free Matmul resources.

It is recommended that the Matmul object be used within a context, but if it is not possible then this method must be called explicitly to ensure that the Matmul object doesn’t hold unnecessary references to operands.

plan( *, stream: AnyStream | int | None = None, ) → None[source]#

Plan the matrix multiplication operation.

Unlike nvmath.linalg.advanced.Matmul.plan(), this method takes no tuning parameters. Its primary function is to find the correct matrix multiplication implementation based on the operands and options provided to the constructor.

Parameters:: stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.
Returns:: None

release_operands()[source]#

Added in version 0.9.0.

This method does two things:

Releases internal references to the user-provided operands, so that this instance no longer contributes to their reference counts.
Frees any internal copies (mirrors) that were created when the user-provided operands reside in a different memory space than the execution (i.e., copies made during construction or within reset_operands() / reset_operands_unchecked()).

This functionality can be useful in memory-constrained scenarios, e.g. where multiple stateful objects need to coexist. Leveraging this functionality, the caller can reduce memory usage while retaining the planned state.

Parameters:: None
Returns:: None

Semantics:

Preserves the planned state of the stateful object.
Repeated calls are safe: if the operands are already released, the call is a no-op and an informational message is logged.
After calling this method, reset_operands() (or reset_operands_unchecked()) must be called to supply new operands before calling execute() again. Failure to do so will result in a runtime error. Device-side copies will be re-allocated as needed.
For cross-space scenarios (e.g. CPU operands with GPU execution, or GPU operands with CPU execution): execution is guaranteed to be always blocking. It is therefore always safe to call this method after calling execute() without additional synchronization.
When the operands are in the same memory space as the execution (e.g. GPU operands with GPU execution): in such case, this method drops this instance’s internal reference to the user-provided operands. If the reference count of the operands reaches zero, their memory may be freed, so particular attention should be paid. The caller is responsible to ensure that if such deallocation happens, it is ordered after pending computation (e.g. by retaining a reference until the computation is complete, or by synchronizing the stream). Failure to do so is analogous to use-after-free.

See Overview, Stateful APIs: Design and Usage Patterns for operand lifecycle and usage patterns, and Stream Semantics for stream ordering rules.

reset_operands( *, a=None, b=None, c=None, alpha=None, beta=None, stream: AnyStream | int | None = None, )[source]#

Reset one or more of the operands held by this Matmul instance.

Changed in version 0.9: All parameters are now keyword-only.

Parameters:

a – A tensor or a sequence of tensors representing the first operand to the matrix multiplication (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
b – A tensor or a sequence of tensors representing the second operand to the matrix multiplication (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
c – (Optional) A tensor or a sequence of tensors representing the operand to add to the matrix multiplication result (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
alpha – The scale factor or a sequence of scale factors for the matrix multiplication term as a real or complex number. The default is \(1.0\).
beta – The scale factor or a sequence of scale factors for the matrix addition term as a real or complex number. A value for beta must be provided if operand c is specified.
stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

Semantics:

Only the operands explicitly passed are updated. At least one operand is required (all of them after release_operands()), otherwise a ValueError is raised.
This method will perform various checks on the new operands to make sure:
- The explicit batch counts, shapes, strides, datatypes match those of the old ones.
- The packages that the operands belong to match those of the old ones.
- If input tensors are on GPU, the device must match.

Examples

>>> import cupy as cp
>>> import nvmath

Create two 3-D float64 ndarrays on the GPU:

>>> M, N, K = 128, 128, 256
>>> a = cp.random.rand(M, K)
>>> b = cp.random.rand(K, N)

Create an matrix multiplication object as a context manager

>>> with nvmath.linalg.Matmul(a, b) as mm:
...     # Plan the operation.
...     mm.plan()
...
...     # Execute the MM to get the first result.
...     r1 = mm.execute()
...
...     # Reset the operands to new CuPy ndarrays.
...     a_new = cp.random.rand(M, K)
...     b_new = cp.random.rand(K, N)
...     mm.reset_operands(a=a_new, b=b_new)
...
...     # Execute to get the new result corresponding to the updated operands.
...     r2 = mm.execute()

Note that if only a subset of operands are reset, the operands that are not reset hold their original values.

With reset_operands(), minimal overhead is achieved as problem specification and planning are only performed once.

For the particular example above, a slower alternative to calling reset_operands() would be to modify the operands in-place (e.g., a[:] = a_new and b[:] = b_new), but this approach is less efficient as it involves copying data rather than just updating pointers.

For more details, please refer to inplace update example.

reset_operands_unchecked( *, a=None, b=None, c=None, alpha=None, beta=None, stream: AnyStream | int | None = None, )[source]#

This method is experimental and potentially subject to future changes.

Added in version 1.0.

This method is a performance-optimized alternative to reset_operands() that eliminates validation and logging overhead, making it ideal for performance-critical loops where operands compatibility is guaranteed by the caller.

This method accepts the same parameters as reset_operands().

Semantics:

The semantics are the same as in reset_operands(), except that this method does not perform any validation (e.g. package match, data type match, etc.) or logging.

When to Use:

Performance-critical loops with repeated executions on different operands
After verifying correctness with reset_operands() during development
When operands compatibility is guaranteed by construction or invariant