direct_solver#

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

Solve the distributed dense linear system \(a @ x = b\) for \(x\). This function-form API is a convenience wrapper around the stateful DirectSolver object and is meant for single use (the user needs to perform just one solve, for example).

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.

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.

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.

See also

DirectSolver, DirectSolverOptions.

Examples

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

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 (chosen here for convenience of exposition; a BlockCyclic distribution is generally preferred for \(a\), as it improves load balancing across processes).

Create the local row slabs on the CPU (cuSOLVERMp requires column-major), and 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]:

>>> 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")
>>> idx = np.arange(local_n)
>>> a[idx, rank * local_n + idx] += n

Solve the system with a single call to direct_solver(). By default the solution \(x\) overwrites b in place and is returned, so x is a NumPy ndarray with the same distribution and local shape as b:

>>> distributions = [Slab.X, Slab.X]
>>> x = direct_solver(a, b, distributions=distributions)
>>> assert x is b

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

Notes

  • This function is a convenience wrapper around DirectSolver and is specifically meant for single use. To amortize the cost of the factorization across multiple right-hand sides, use the stateful DirectSolver API.

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

On this page