warp.optim.linear.FSAI#

class warp.optim.linear.FSAI( A, max_row_size=8, kap_tolerance=1.0e-3, pivot_floor=None, factor_dtype=None, reuse_pattern=False, max_step_size=1, regularize_refit=False, validate_topology=True, defer_pattern=False, )[source]#

Build an adaptive factorized sparse approximate inverse of an SPD BSR matrix.

This warp.optim.linear.LinearOperator applies G.T @ G and can be passed as M to warp.optim.linear.cg() or warp.optim.linear.cr(). Rows of the lower-triangular factor G are built independently on the matrix device. Square blocks are expanded into scalar entries, so this is scalar FSAI even when the input uses block storage. No host numerical backend is used.

Parameters:

A (BsrMatrix) – Square warp.sparse.BsrMatrix with square blocks and float16, float32 or float64 scalar entries. Both triangles must be stored. Symmetry and positive definiteness are caller preconditions, not checked globally. Stored columns must be sorted and unique within each row, as produced by the standard BSR construction functions.
max_row_size (int) – Maximum scalar entries per factor row, including its diagonal, in [1, 64]. One entry gives Jacobi preconditioning. Larger supports use quadratic per-thread setup workspace and increase setup and apply costs.
kap_tolerance (float) – Relative local energy improvement below which support growth stops, in [0, 1). Zero disables this stopping test. It is checked after each batch, so changing max_step_size may require retuning it.
pivot_floor (float | None) – Smallest accepted pivot of the diagonally equilibrated local Cholesky factorization, in (0, 1). None selects 1e-3 for float16, 1e-6 for float32 and 1e-12 for float64. Setup truncates a row before a rejected pivot.
factor_dtype (type | None) – Factor storage type: the matrix scalar type, or warp.float32 to compress a float64 factor. None uses the matrix scalar type. Products accumulate in the matrix scalar type; the transpose is formed from the same rounded factor, preserving the Gram form.
reuse_pattern (bool) – Retain a plan for update() on the original factor supports. Requires compact input storage; use warp.sparse.bsr_compress() with prune_numerical_zeros=False to compact a padded matrix first. Ordinary construction accepts either.
max_step_size (int) – Maximum distinct largest-residual frontier entries selected per growth step, in [1, 64]. Batching reduces searches but can change the selected supports. A short batch does not exhaust the row-size budget.
regularize_refit (bool) – With reuse_pattern, make update() refit every row it can instead of rejecting the whole factor when one row fails. A row whose pivot falls to pivot_floor or below is refit with the pivot raised to the floor and counted as regularized. A row with a nonfinite pivot, an invalid diagonal, or a nonpositive result is replaced by its scaled diagonal and counted as failed. Only a topology mismatch rejects the update. Inspect update_status to see affected rows.
validate_topology (bool) – With reuse_pattern, have update() compare the matrix topology with the original on the device and reject a mismatch. Pass False to skip the comparison when the caller guarantees an unchanged topology; a changed topology then gives undefined results and may read out of bounds.
defer_pattern (bool) – Construct from the topology of A alone, without reading its values, and select the supports on the device during the first successful update(). Requires reuse_pattern. Until then the preconditioner is the identity. Every kernel of that selection checks pattern_ready and does nothing once the pattern exists, so a CUDA graph that captured the first update selects the pattern on its first replay and only refits on later replays. G is then stored with padded rows (row_counts), GT has capacity for max_row_size entries per row. A row with an unusable diagonal keeps a diagonal-only support, which is reported in update_status, if regularize_refit is set; otherwise the update is rejected and the next one selects again.

G#: Canonical scalar BSR factor with finite entries and a positive diagonal.

GT#: Stored transpose of G, including any storage rounding.

truncated_rows#: Number of setup rows stopped by the pivot safeguard, or None with defer_pattern (see pattern_status).

pattern_ready#: None without defer_pattern. Otherwise an int array on the matrix device holding 1 once the supports have been selected and published.

pattern_status#: None without defer_pattern. Otherwise an int array on the matrix device with the counts from the most recent support selection: the rows stopped by the pivot safeguard at index 1 and the rows whose factor lost finite entries or a positive diagonal at index 2. Index 0 is unused.

source#: Matrix used by the most recent successful construction or update.

update_status#: None without reuse_pattern. Otherwise an int array on the matrix device with four counts from the most recent update: the topology mismatches against the original matrix, the invalid diagonals, the local factorizations that failed, and the rows regularized (always zero unless regularize_refit). By default the new factors were applied only if the first three are zero; with regularize_refit they were applied unless the topology mismatched.

Note

Construction synchronizes and cannot be captured, but with defer_pattern it needs only the topology of the matrix. update() can be captured in a CUDA graph (see its documentation), and application is CUDA graph-capturable after warming up the kernels, including inside the conditional graph loop of the iterative solvers. Applications of one instance on different streams are ordered against each other and never overlap. Do not overlap an update with application, and do not replay one captured graph that applies the instance concurrently with another use of the same instance. Treat G and GT as read-only. Rebuild or call update() after changing the input values. This preconditioner does not provide automatic differentiation.

Raises:

TypeError – The matrix scalar type is unsupported.
ValueError – The matrix shape or an option is invalid, a diagonal is missing, nonpositive or nonfinite, or the factor cannot be represented by the selected storage type with finite entries and a positive diagonal.

Parameters:

A (BsrMatrix)
max_row_size (int)
kap_tolerance (float)
pivot_floor (float | None)
factor_dtype (type | None)
reuse_pattern (bool)
max_step_size (int)
regularize_refit (bool)
validate_topology (bool)
defer_pattern (bool)

Example

from warp.optim.linear import FSAI, cg

M = FSAI(A, max_row_size=16, reuse_pattern=True)
cg(A, b, x, M=M)
# After updating A.values without changing A's sparsity:
M.update()
cg(A, b, x, M=M)

Methods

update([A])

Refit on the original factor supports, preserving captured apply buffers.

Attributes

`batch_count`	Number of independent subproblems.
`batch_offsets`	Array partitioning the active scalar degrees of freedom, or `None`.
`device`	Device on which computations involving the operator are performed.
`dtype`	Type of the operator elements, such as a scalar type or a matrix block type.
`matvec`	Matrix-vector multiplication function.
`max_batch_length`	Upper bound on the number of scalar degrees of freedom in any subproblem, or `None`.
`scalar_type`	Scalar component type of the operator's elements.
`shape`	Number of rows and columns of the operator.

update(A=None)[source]#

Refit on the original factor supports, preserving captured apply buffers.

Requires reuse_pattern=True and identical compact BSR topology/storage. This is not a new adaptive pattern search, except that with defer_pattern the first successful update selects the supports. Rebuild when convergence worsens. Failed validation leaves the previous factors intact.

Outside graph capture, the update synchronizes and raises on failure, except that with both regularize_refit and validate_topology=False nothing can raise, so it does not synchronize. During CUDA graph capture, the update performs no host reads and cannot raise for numerical failures; after each replay, inspect update_status. Run one eager update first so the kernels are loaded before capture, and have the matrix keep the same arrays so a replay reads its new values.

Parameters:: A (BsrMatrix | None) – Matrix with new numerical values. None reuses source.
Returns:: This preconditioner with updated numerical factors.
Raises:: ValueError – No reuse plan exists, matrix metadata or array sizes changed, or, when not capturing, the topology changed or a local factorization or storage conversion failed. Unlike initial construction, refits reject a bad pivot without truncating the fixed support. Rebuild to select new supports.
Return type:: FSAI