DirectSolver#

Create a stateful object that encapsulates the specified dense direct linear solver computations and required resources. This object ensures the validity of resources during use and releases them when they are no longer needed to prevent misuse.

This object encompasses all functionality of the function-form API direct_solver(), which is a convenience wrapper around it. The stateful object also allows amortization of preparatory costs (for example plan() and an LU factorize()) when the same size and layout are reused, including after reset_operands() to update a/b.

Using the stateful object typically involves the following steps:

Problem specification: Initialize the object with a, b, and optional options / execution / stream.
Preparation: plan() in preparation for factorization (see Semantics).
Execution: factorize(), then solve() to obtain x (or solve() again with the same a and an updated b after reset_operands() when the matrix or RHS values change without recreating the solver).
Resource management: Release resources with free() or by using a with block.

More detail on what each phase is doing (planning, factorization, and solve) can be seen in the logs by passing a logging.Logger in DirectSolverOptions, or by configuring the root logger that DirectSolverOptions uses by default, for example:

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

Parameters:

a – The dense operand (or sequence of operands) representing the left-hand side (LHS) of the system of equations: a square matrix a (or batch of square (n, n) matrices) for a @ x = b. The LHS may be a (sequence of) numpy.ndarray, cupy.ndarray, and torch.Tensor. For batched input, a can also be a single tensor of shape (*batch_shape, n, n) (implicit batching) or a list or tuple of matrices of the same shape. Refer to the Semantics section for details.
b – The ndarray/tensor or (sequence of ndarray/tensors) representing the dense right-hand side (RHS) of the system of equations. The RHS operand may be a (sequence of) numpy.ndarray, cupy.ndarray, and torch.Tensor. Unbatched shapes are (n,) or (n, nrhs); with batching parallel to a, use a list / tuple of 1D/2D arrays (explicit) or a single tensor (*batch_shape, n, nrhs) (implicit). The package and device must match a. Refer to the Semantics section for details.
options – Specify options for the direct linear 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.
execution – Specify execution space options for the direct solver as an ExecutionCUDA object. Alternatively, a string ("cuda") or a dict with a "name" of "cuda" and optional device_id relevant to the given execution space. The default execution space is "cuda" and the corresponding ExecutionCUDA will be default-constructed on device 0. CPU operands are executed on the selected CUDA device as needed.
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 dense direct linear solver solves \(a @ x = b\) for x given the left-hand side (LHS) a and the right-hand side (RHS) b (square dense \(a\), dense b), using an LU factorization and triangular solves on the target device.

Non-batched: a is a square (n, n) matrix; b is a vector (n,) or a matrix (n, nrhs) (one or more right-hand sides). The solve runs on the selected CUDA device using cuSOLVER.
Batched (explicit or implicit): the solve uses cuBLAS batched APIs, so in one call every batch entry shares the same system size \(n\) (all (n, n) a). For b, explicit batching uses per-item shapes (n,) or (n, m); implicit batching uses (*batch_shape, n, nrhs). If a and b are implicitly batched, they must share the same leading batch_shape where applicable. The LHS and RHS batching forms may differ (for example, a sequence of a with an implicitly batched b, or the reverse) as long as the number of batch entries is the same for a and b.

See also

nvmath.sparse.advanced.DirectSolver for \(a\) as a sparse CSR matrix (a different feature set and backend).

Note

Currently supported dtypes are: float32, float64, complex64, complex128 for both a and b.
If users provide options.handle, it is used only for non-batched execution (cuSOLVER path). In batched execution (cuBLAS path), the provided handle is ignored and an internal cached cuBLAS handle is used.
When inplace_a=True or inplace_b=True for batched execution, users are responsible for ensuring that the LHS and RHS batch entries do not overlap in memory.

Examples

Non-batched numpy.ndarray operands: build a well-conditioned \(n \times n\) system and solve with plan(), factorize(), and solve() inside a context manager so that free() is called automatically:

>>> import numpy as np
>>> import nvmath.linalg as la
>>>
>>> n = 4
>>> rng = np.random.default_rng(0)
>>> a = rng.standard_normal((n, n), dtype=np.float64)
>>> a = a + n * np.eye(n, dtype=np.float64)
>>> b = np.ones(n, dtype=np.float64)
>>>
>>> with la.DirectSolver(a, b) as solver:
...     solver.plan()
...     solver.factorize()
...     x = solver.solve()
>>> x.shape
(4,)

The same call pattern applies to cupy.ndarray and torch.Tensor on the default CUDA device when the operands already live on the GPU, and to batching; see Semantics and the example-tree link below. Further worked examples (implicit batching, reset_operands(), streams) are in generic direct solver examples on GitHub.

Methods

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

Factorize the system of equations. This performs a numerical LU decomposition of the left-hand side a for the current problem. Call plan() first.

Each call factorizes the LHS a currently held by the solver. To change the LHS 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 LHS 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()[source]#

Free DirectSolver resources.

It is recommended that the object be used as a context manager; if that is not possible, call this method explicitly when the object is no longer needed.

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

Plan the dense direct solve. This is the preparation step before numerical factorization (any sizing or setup required for the current operands on the selected device). After planning, the object is ready for factorize().

In the class-based workflow, call plan() once before calling factorize() and solve().

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

Replace a and/or b with operands that match the same problem shape, batching, device, package, and dtype as the original.

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:

Only the operands explicitly passed are updated. At least one operand is required (all of them after release_operands()), otherwise a ValueError is raised.
Supplying a new a invalidates the current factorization; call factorize() again before the next solve().

Note

When inplace_a=True or inplace_b=True and the input operands reside on the GPU, replacement operands must have the same layout as the original operands. CPU operands are copied through internal GPU buffers, so this layout restriction does not apply to CPU inputs.

reset_operands_unchecked( *, a=None, b=None, stream: AnyStream | int | 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 a/b compatibility is guaranteed by the caller.

This method accepts the same parameters as reset_operands().

Parameters:

a – New left-hand side, or None to keep the previous one. The caller must guarantee the same invariants reset_operands() would enforce; this method does not validate.
b – New right-hand side, or None to keep the previous one. Same invariants as a where applicable.
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

Semantics:

The semantics are the same as in reset_operands(), except that this method does not perform any validation (e.g. batching metadata, n, n_rhs, device, and dtype) or logging.

When to Use:

Performance-critical loops with repeated solve() or operand updates when compatibility is already guaranteed
After verifying correctness with reset_operands() during development
When a/b compatibility is guaranteed by construction or invariant

Note

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

Solve the factorized system of equations, producing the solution x for \(a @ x = b\) using the factor from the most recent factorize() and the current right-hand side b.

Each call uses the RHS b currently held by the solver. To change the RHS 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 RHS 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:: The result of the specified direct linear solve, x, which has the same batching structure as the RHS b (a single array/tensor, or a sequence of ndarray/tensors, each the same shape as the corresponding tensor in b), a dtype matching a and b, and device and package placement consistent with the chosen execution and operands.