Matmul#

class nvmath.distributed.linalg.advanced.Matmul(
a,
b,
/,
c=None,
*,
distributions: Sequence[Distribution],
alpha=None,
beta=None,
qualifiers=None,
options=None,
stream: AnyStream | int | None = None,
)[source]#

Create a stateful object encapsulating the specified distributed 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:

  1. Problem Specification: Initialize the object with a defined operation and options.

  2. Preparation: Use plan() to determine the best algorithmic implementation for this specific matrix multiplication operation.

  3. Execution: Perform the matrix multiplication computation with execute().

  4. Resource Management: Ensure all resources are released either by explicitly calling free() or by managing the stateful object within a context manager.

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 distributed tensor representing the first operand to the matrix multiplication (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.

  • b – A distributed tensor 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 distributed tensor representing the operand to add to the matrix multiplication result (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.

  • distributions – Sequence specifying the distribution across processes of matrices A, B and C/D. The distribution needs to be BlockCyclic or compatible.

  • alpha – The scale factor for the matrix multiplication term as a real or complex number. The default is \(1.0\).

  • beta – The scale factor for the matrix addition term as a real or complex number. A value for beta must be provided if operand c is specified.

  • qualifiers – Specify the matrix qualifiers as a numpy.ndarray of matrix_qualifiers_dtype objects of length 3 corresponding to the operands a, b, and c. 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.

  • 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.

Semantics:

The semantics of the matrix multiplication follows numpy.matmul() semantics, with some restrictions on broadcasting.

See also

plan(), reset_operands(), execute()

Examples

>>> import numpy as np
>>> import nvmath.distributed
>>> from nvmath.distributed.distribution import Slab
>>> from nvmath.distributed.linalg.advanced import matrix_qualifiers_dtype

Get MPI communicator used to initialize nvmath.distributed (for information on initializing nvmath.distributed, you can refer to the documentation or to the Matmul examples in nvmath/examples/distributed/linalg/advanced):

>>> comm = nvmath.distributed.get_context().communicator

Get my process rank:

>>> rank = comm.Get_rank()

Create two 2-D float64 ndarrays on the CPU (using Slab distributions to distribute the matrices across processes):

>>> M, N, K = 1024, 1024, 1024
>>> a_shape = Slab.X.shape(rank, (K, M))
>>> b_shape = Slab.X.shape(rank, (K, N))
>>> a = np.asfortranarray(np.random.rand(*a_shape))
>>> b = np.asfortranarray(np.random.rand(*b_shape))

We will define a matrix multiplication operation followed by an AllReduce epilog using the specialized matrix multiplication interface.

Create a Matmul object encapsulating the problem specification above:

>>> qualifiers = np.zeros((3,), dtype=matrix_qualifiers_dtype)
>>> qualifiers[0]["is_transpose"] = True  # a is transposed
>>> distributions = [Slab.X, Slab.X, Slab.Y]
>>> mm = nvmath.distributed.linalg.advanced.Matmul(
...     a, b, distributions=distributions, qualifiers=qualifiers
... )

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

Next, plan the operation. The epilog is specified, and optionally, preferences can be specified for planning:

>>> epilog = nvmath.distributed.linalg.advanced.MatmulEpilog.ALLREDUCE
>>> mm.plan(epilog=epilog)

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

>>> r1 = mm.execute()

Finally, free the object’s resources. To avoid having to explicitly making this call, it’s recommended to use the Matmul object as a context manager as shown below, if possible.

>>> mm.free()

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.

>>> device_id = nvmath.distributed.get_context().device_id
>>> import cupy as cp
>>> with cp.cuda.Device(device_id):
...     a = cp.asfortranarray(cp.random.rand(*a_shape))
...     b = cp.asfortranarray(cp.random.rand(*b_shape))

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

>>> with nvmath.distributed.linalg.advanced.Matmul(
...     a, b, distributions=distributions, qualifiers=qualifiers
... ) as mm:
...     mm.plan(epilog=epilog)
...
...     # Execute the operation to get the first result.
...     r1 = mm.execute()
...
...     # Update operands A and B in-place (see reset_operands() for an
...     # alternative).
...     with cp.cuda.Device(device_id):
...         a[:] = cp.random.rand(*a_shape)
...         b[:] = cp.random.rand(*b_shape)
...
...     # 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 nvmath/examples/distributed/linalg/advanced/matmul directory.

Methods

__init__(
a,
b,
/,
c=None,
*,
distributions: Sequence[Distribution],
alpha=None,
beta=None,
qualifiers=None,
options=None,
stream: AnyStream | int | None = None,
)[source]#
execute(
*,
release_workspace=False,
stream: AnyStream | int | None = None,
)[source]#

Execute a planned distributed matrix multiplication.

Parameters:

  • release_workspace – A value of True specifies that the stateful object should release workspace memory back to the package memory pool on function return, while a value of False specifies that the object should retain the memory. This option may be set to True if the application performs other operations that consume a lot of memory between successive calls to the (same or different) execute() API, but incurs a small overhead due to obtaining and releasing workspace memory from and to the package memory pool on every call. The default is False.

  • 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.

Returns:

The result of the specified matrix multiplication (epilog applied), which remains on the same device and belongs to the same package as the input operands.

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 matrix multiplication resources (especially internal library objects) are properly cleaned up.

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

Plan the matrix multiplication operation, considering the epilog (if provided).

Parameters:

  • preferences – This parameter specifies the preferences for planning as a MatmulPlanPreferences object. Alternatively, a dictionary containing the parameters for the MatmulPlanPreferences constructor can also be provided. If not specified, the value will be set to the default-constructed MatmulPlanPreferences object.

  • epilog – Specify an epilog \(F\) as an object of type MatmulEpilog to apply to the result of the matrix multiplication: \(F(\alpha A @ B + \beta C\)). The default is no epilog. See cuBLASMp documentation for the list of available epilogs.

  • epilog_inputs – Specify the additional inputs needed for the selected epilog as a dictionary, where the key is the epilog input name and the value is the epilog input. The epilog input must be a tensor with the same package and in the same memory space as the operands (see the constructor for more information on the operands). If the required epilog inputs are not provided, an exception is raised that lists the required epilog inputs. Some epilog inputs are generated by other epilogs. For example, the epilog input for MatmulEpilog.DRELU is generated by matrix multiplication with the same operands using MatmulEpilog.RELU_AUX.

  • 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 Matmul for an example, and further examples can be found in the nvmath/examples/distributed/linalg/advanced/matmul directory.

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

Reset the operands held by this Matmul instance.

This method has two use cases:

  1. it can be used to provide new operands for execution when the original operands are on the CPU

  2. it can be used to release the internal reference to the previous operands and make their memory available for other use by passing None for all arguments. In this case, this method must be called again to provide the desired operands before another call to execute().

This method is not needed when the operands reside on the GPU and in-place operations are used to update the operand values.

This method will perform various checks on the new operands to make sure:

  • The distributions, 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.

Parameters:

  • a – A distributed tensor representing the first operand to the matrix multiplication (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.

  • b – A distributed tensor 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 distributed tensor 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 for the matrix multiplication term as a real or complex number. The default is \(1.0\).

  • beta – The scale factor for the matrix addition term as a real or complex number. A value for beta must be provided if operand c is specified.

  • epilog_inputs – Specify the additional inputs needed for the selected epilog as a dictionary, where the key is the epilog input name and the value is the epilog input. The epilog input must be a tensor with the same package and in the same memory space as the operands (see the constructor for more information on the operands). If the required epilog inputs are not provided, an exception is raised that lists the required epilog inputs. Some epilog inputs are generated by other epilogs. For example, the epilog input for MatmulEpilog.DRELU is generated by matrix multiplication with the same operands using MatmulEpilog.RELU_AUX.

  • 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.

Examples

>>> import cupy as cp
>>> import nvmath.distributed
>>> from nvmath.distributed.distribution import Slab

Get MPI communicator used to initialize nvmath.distributed (for information on initializing nvmath.distributed, you can refer to the documentation or to the Matmul examples in nvmath/examples/distributed/linalg/advanced):

>>> comm = nvmath.distributed.get_context().communicator

Get my process rank:

>>> rank = comm.Get_rank()

Create two 3-D float64 ndarrays on the GPU (using Slab distributions to distribute the matrices across processes):

>>> M, N, K = 128, 128, 256
>>> a_shape = Slab.X.shape(rank, (M, K))
>>> b_shape = Slab.Y.shape(rank, (K, N))
>>> device_id = nvmath.distributed.get_context().device_id
>>> with cp.cuda.Device(device_id):
...     a = cp.asfortranarray(cp.random.rand(*a_shape))
...     b = cp.asfortranarray(cp.random.rand(*b_shape))

Create an matrix multiplication object as a context manager

>>> d = [Slab.X, Slab.Y, Slab.X]
>>> with nvmath.distributed.linalg.advanced.Matmul(a, b, distributions=d) 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.
...     with cp.cuda.Device(device_id):
...         c = cp.asfortranarray(cp.random.rand(*a_shape))
...         d = cp.asfortranarray(cp.random.rand(*b_shape))
...     mm.reset_operands(c, d)
...
...     # 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, explicitly calling reset_operands() is equivalent to updating the operands in-place, i.e, replacing mm.reset_operand(c, d) with a[:]=c and b[:]=d. Note that updating the operand in-place should be adopted with caution as it can only yield the expected result under the additional constraint below:

  • The operand is on the GPU (more precisely, the operand memory space should be accessible from the execution space).

For more details, please refer to inplace update example.