The nvmath.device module is experimental and potentially subject to future changes.

CholeskySolver#

class nvmath.device.CholeskySolver( size: Sequence[int], precision: type[floating], execution: str, fill_mode: str, *, sm=None, arrangement: Sequence[str] | None = None, batches_per_block: int | Literal['suggested'] | None = None, data_type: str | None = None, leading_dimensions: Sequence[int] | None = None, block_dim: Sequence[int] | Literal['suggested'] | None = None, )[source]#

A class that encapsulates Cholesky factorization and solve device functions for Hermitian positive definite matrices.

Available operations:

factorize: Computes the Cholesky factorization A = L @ L^H (lower) or A = U^H @ U (upper), where L is a lower triangular matrix and U is an upper triangular matrix. The choice depends on the fill_mode parameter.
solve: Solves the system Ax = B using a previously computed Cholesky factorization

Memory Layout Requirements:

Matrices must be stored in shared memory according to their arrangement and leading dimension (ld):

For matrix A (M x N):

Column-major arrangement: Matrix shape (batches_per_block, M, N) with strides (lda * N, 1, lda)
Row-major arrangement: Matrix shape (batches_per_block, M, N) with strides (lda * M, lda, 1)

For matrix B (N x K):

Column-major arrangement: Matrix shape (batches_per_block, N, K) with strides (ldb * K, 1, ldb)
Row-major arrangement: Matrix shape (batches_per_block, N, K) with strides (ldb * N, ldb, 1)

Parameters:

size (Sequence[int]) – Problem size specified as a sequence of 1 to 3 elements: (M,) (treated as (M, M, 1)), (M, N) (treated as (M, N, 1)), or (M, N, K). M represents the dimension of the square matrix A (M x M) used in factorization, N must be equal to M. K represents the number of columns in the right-hand side matrix B (dimensions M x K) for the solve operation.
precision (type[np.floating]) – The computation precision specified as a numpy float dtype. Currently supports: numpy.float32, numpy.float64.
execution (str) – A string specifying the execution method. Supported values: 'Block'.
fill_mode (str) – Indicates which part of matrix A is filled and should be used by function. Can be one of: 'upper', 'lower'.
sm (ComputeCapability) – Target mathdx compute-capability.
arrangement (Sequence[str], optional) – Storage layout for matrices A and B, specified as a sequence of 1 or 2 elements (arr_A, arr_B). Each element can be one of: 'col_major', 'row_major'. Defaults to ("col_major", "col_major").
batches_per_block (int | Literal["suggested"], optional) – Number of batches to compute in parallel in a single CUDA block. Can be a positive integer or the string 'suggested' for automatic selection of an optimal value. We recommend using 1 for matrix A size larger than or equal to 16 x 16, and using 'suggested' for smaller sizes to achieve optimal performance. Defaults to 1.
data_type (str, optional) – The data type of the input matrices, can be one of: 'real', 'complex'. Defaults to 'real'.
leading_dimensions (Sequence[int], optional) – The leading dimensions for input matrices A and B, specified as a sequence of 1 or 2 elements (lda, ldb) or None. If not provided, it will be automatically deduced from size and arrangement. Note: When provided in the constructor, leading dimensions are set at compile-time. To use runtime leading dimensions (avoiding recompilation for different leading dimensions), provide the leading dimension parameters directly to the device methods instead.
block_dim (Sequence[int] | Literal["suggested"], optional) – The block dimension for launching the CUDA kernel, specified as a 1 to 3 integer sequence (x, y, z) where missing dimensions are assumed to be 1. Can be a sequence of 1 to 3 positive integers, the string 'suggested' for optimal value selection, or None for the default value.