DirectSolver#

class nvmath.distributed.linalg.DirectSolver( a: AnyTensor, b: AnyTensor, /, *, distributions: Sequence[Distribution], options: DirectSolverOptions | dict[str, Any] | None = None, stream: AnyStream | int | None = None, )[source]#

Create a stateful object that encapsulates the specified distributed direct solver computation for the dense linear system \(a @ x = b\) and the required resources to perform it.

A stateful object can be used to amortize the cost of preparation (planning and factorization) across multiple solves with different right-hand sides, or with updated operands of the same problem specification (see reset_operands()). The function-form API direct_solver() is a convenient alternative for single use (the user needs to perform just one solve, for example), in which case there is no possibility of amortizing preparatory costs.

Using the stateful object typically involves the following steps:

Problem Specification: Initialize the object with the operands, their distributions, and options.
Preparation: Use plan() to query and size the solver workspaces for this specific problem.
Execution: factorize() the matrix and solve() the system.
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 is happening in the various phases described above can be obtained by passing in a logging.Logger object to DirectSolverOptions 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",
... )

Parameters:

a – A distributed tensor representing the left-hand side operand \(a\) of the linear system \(a @ x = b\) (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
b – A distributed tensor representing the right-hand side \(b\) of the linear system \(a @ x = b\) (see Semantics). The currently supported types are numpy.ndarray, cupy.ndarray, and torch.Tensor.
distributions – A sequence [distribution_for_a, distribution_for_b] specifying how \(a\) and \(b\) are partitioned across processes. See Semantics and Requirements for the supported distribution types and the constraints they must satisfy (in particular, the row block size of \(b\) must match that of \(a\)).
options – Specify options for the direct solver as a DirectSolverOptions object. Alternatively, a dict containing the parameters for the DirectSolverOptions constructor can also be provided. If not specified, the value will be set to the default-constructed DirectSolverOptions 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. See Stream Semantics for more details on stream handling.

Semantics:

The distributed direct solver solves the dense linear system \(a @ x = b\) for \(x\), where \(a\) is a square (n, n) matrix and \(b\) has shape (n, nrhs) or (n,), where nrhs is the number of RHS vectors. It performs an LU factorization with partial pivoting (DirectSolver.factorize()) followed by triangular solves (DirectSolver.solve()). The computation always runs on the GPU; CPU operands are copied to the device and back automatically.

On each process, a and b hold only the portions of the global matrices \(a\) and \(b\) that are assigned to that process. The corresponding entries of distributions describe how \(a\) and \(b\) are partitioned across processes.

The solution \(x\) has the same local shape, memory space, and package as b and follows b’s distribution. By default, DirectSolver.factorize() overwrites a with its LU factors and DirectSolver.solve() overwrites b with \(x\) and returns b; this in-place behavior is controlled by the inplace_a and inplace_b attributes of DirectSolverOptions (both True by default). If \(a\) is singular, DirectSolver.factorize() raises.

Requirements:

Operands.

The global number of columns of a must match the global number of rows of b (both equal \(n\), the order of the square system).
a and b must belong to the same package and have the same dtype, one of float32, float64, complex64, or complex128.
a and b must reside in the same memory space on each process, and GPU operands must be on the device bound to the distributed runtime on that process.
The local arrays must be column-major; a padded leading dimension (a column-major slice of a larger buffer) is allowed.

Distributions.

distributions must contain exactly two entries, [distribution_for_a, distribution_for_b], respectively describing how a and b are partitioned across processes.

The simplest, but not most performant, choice is to give both a and b the same row-wise Slab distribution (Slab.X), so that each process owns an equal, contiguous block of rows of the global operands. The only requirement here is that the global number of rows be divisible by the number of processes (a Slab with unequal partition sizes is not supported, because it cannot be expressed as a block-cyclic distribution); a and b are then partitioned consistently and all the other requirements below are satisfied automatically, so none of them need to be considered.

In general, each distribution entry may be a Slab (with uniform partition sizes), BlockNonCyclic, or BlockCyclic distribution, and a and b may use different distributions. This flexibility comes with stricter requirements:

The process grid must place its first process at (0, 0) and use a COL_MAJOR or ROW_MAJOR layout.
The row block size of b must equal that of a, while column block sizes are independent.
A 1-D b requires b’s distribution to be 1-D (row-only): a 2-D distribution is rejected, because the single column cannot be split across process columns.
For a 2-D b, b’s distribution may be either 1-D (row-only) or 2-D, partitioning both rows and columns.

Examples

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

Get the process group used to initialize nvmath.distributed (for information on initializing nvmath.distributed, you can refer to the documentation or to the direct solver examples in nvmath/examples/distributed/linalg/generic/direct_solver):

>>> process_group = nvmath.distributed.get_context().process_group

Get my process rank and the total number of processes:

>>> rank = process_group.rank
>>> nranks = process_group.nranks

We will solve the dense square system \(a @ x = b\), where \(a\) has shape (n, n) and \(b\) has shape (n, nrhs). Both operands use a Slab distribution, so each process owns a contiguous block of rows of the global matrices.

Note

The Slab distribution is used here for convenience of exposition and is not necessarily the most performant choice. For \(a\), a BlockCyclic distribution is generally preferred, as it improves load balancing across processes.

Create the local row slabs on the CPU (cuSOLVERMp requires column-major):

>>> n, nrhs = 256, 8
>>> local_n = n // nranks  # assume n is divisible by the process count
>>> rng = np.random.default_rng(rank)
>>> a = rng.random((local_n, n)).astype(np.float64, order="F")
>>> b = rng.random((local_n, nrhs)).astype(np.float64, order="F")

Make the global \(a\) diagonally dominant so the system is well-conditioned. The diagonal entry a[i, i] of this rank’s row band is at local position [k, rank * local_n + k]:

>>> idx = np.arange(local_n)
>>> a[idx, rank * local_n + idx] += n

Create a DirectSolver object encapsulating the problem specification above:

>>> distributions = [Slab.X, Slab.X]
>>> solver = DirectSolver(a, b, distributions=distributions)

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

Next, plan the operation:

>>> solver.plan()

Factorize the matrix, then solve the system. By default, the solution \(x\) overwrites b and is returned:

>>> solver.factorize()
>>> x = solver.solve()

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

>>> solver.free()

A benefit of the stateful API is that an expensive factorization can be reused across multiple right-hand sides: factorize once, then solve repeatedly, resetting only b between solves.

Further examples can be found in the nvmath/examples/distributed/linalg/generic/direct_solver directory.

Methods

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

Factorize the matrix \(a\) (LU factorization with partial pivoting).

Each call factorizes the left-hand side a currently held by the solver. To change the left-hand side values between calls, call reset_operands() (or reset_operands_unchecked()) with the new a. Alternatively, when inplace_a=True and a is accessible from the execution space, users may directly modify a in place, for example with a[:] = a_new. In that case, users are responsible for ensuring that a contains the intended left-hand side values before each call, since factorize() overwrites that operand with its LU factors.

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.

free() → None[source]#

Free the DirectSolver resources.

It is recommended that the DirectSolver object be used within a context manager, but if that is not possible this method must be called explicitly to ensure the solver’s resources are properly cleaned up.

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

Plan the distributed direct solver. This queries cuSOLVERMp for the host and device workspace sizes needed by the factorization and the solve, and sizes the solver’s internal workspaces accordingly. Planning is idempotent, so calling it again after a successful plan is a no-op.

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.

See also

factorize(), solve().

release_operands() → None[source]#

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 factorize() / solve() 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 factorize() / solve() 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: Any = None, b: Any = None, stream: AnyStream | int | None = None, ) → None[source]#

Reset one or both operands held by this DirectSolver instance.

Parameters:

a – New left-hand side, or None to keep the previous left-hand side.
b – New right-hand side, or None to keep the previous right-hand side.
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:

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 local shapes, strides, and datatypes match those of the old ones.
- The packages that the operands belong to match those of the old ones.
- The device must match that of the old ones.
The distribution of each operand is fixed at construction and cannot be changed by reset_operands(); the new local shapes and strides need to be consistent with that distribution.
Resetting a invalidates the current factorization, so factorize() must be called again before the next solve().

reset_operands_unchecked( *, a: Any = None, b: Any = None, stream: AnyStream | int | None = None, ) → None[source]#

This method is experimental and potentially subject to future changes.

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

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

Solve the factorized system \(a @ x = b\) for \(x\) using the factors from the most recent factorize() and the current right-hand side b.

Each call uses the right-hand side b currently held by the solver. To change the right-hand side values between calls, call reset_operands() (or reset_operands_unchecked()) with the new b. Alternatively, when inplace_b=True and b is accessible from the execution space, users may directly modify b in place, for example with b[:] = b_new. In that case, users are responsible for ensuring that b contains the intended right-hand side values before each call, since solve() overwrites that operand with the solution.

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:: A distributed tensor representing the solution \(x\) of the linear system \(a @ x = b\) (see Semantics). It has the same distribution, local shape, device, and package as the right-hand side b. When the inplace_b attribute of DirectSolverOptions is True (the default), \(x\) is written into b in place and the returned object is b itself.