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 symmetric 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'.

  • sm (ComputeCapability) – Target mathdx compute-capability.

  • fill_mode (str) – Indicates which part of matrix A is filled and should be used by function. Can be one of: 'upper', 'lower'.

  • arrangement (Sequence[str], optional) – Storage layout for matrices A and B, specified as a sequence of 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 non-zero 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 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.

See also

For further details, please refer to the cuSOLVERDx documentation:

Attributes

a_arrangement#
a_shape#
arrangement#
b_arrangement#
b_shape#
batches_per_block#
block_dim#
block_size#
data_type#
execution#
fill_mode#
info_shape#
info_strides#
info_type#
k#
lda#
ldb#
leading_dimensions#
m#
n#
precision#
size#
sm#
value_type#

Methods

a_size(*, lda: int | None = None) int[source]#
a_strides(
*,
lda: int | None = None,
) tuple[int, int, int][source]#
b_size(*, ldb: int | None = None) int[source]#
b_strides(
*,
ldb: int | None = None,
) tuple[int, int, int][source]#
factorize(a, info, lda=None) None[source]#

Computes the Cholesky factorization of a symmetric positive definite matrix A.

This device function computes A = L @ L^H (if fill_mode = 'lower') or A = U^H @ U (if fill_mode = 'upper'). Uses cuSOLVERDx 'potrf'.

If lda is provided, uses runtime version with the specified leading dimension. If lda is not provided (None), uses compile-time version with default or constructor-provided leading dimensions.

For more details, see: get_started/functions/potrf.html

Parameters:
  • a – Pointer to an array in shared memory, storing the matrix according to the specified arrangement and leading dimension (see __init__()). On entry, contains the symmetric positive definite matrix. On exit, contains the triangular factor L (lower) or U (upper).

  • info – Pointer to a 1D array of int32. On exit, info[batch_id] = 0 indicates success for that batch, info[batch_id] != 0 indicates the matrix is not positive definite.

  • lda – Optional runtime leading dimension of matrix A. If not specified, the compile-time lda is used.

solve(a, b, lda=None, ldb=None) None[source]#

Solves a system of linear equations Ax = B using the Cholesky factorization.

This device function uses the previously computed factorization A = L @ L^H (lower) or A = U^H @ U (upper) to solve the system. Uses cuSOLVERDx 'potrs'.

If lda and ldb are provided, uses runtime version with the specified leading dimensions. If not provided (None), uses compile-time version with default or constructor-provided leading dimensions.

For more details, see: get_started/functions/potrs.html

Parameters:
  • a – Pointer to an array in shared memory, storing the triangular factor L (lower) or U (upper) from the Cholesky factorization, according to the specified arrangement and leading dimension (see __init__()).

  • b – Pointer to an array in shared memory, storing the matrix according to the specified arrangement and leading dimension (see __init__()). The matrix is overwritten in place with the solution matrix x.

  • lda – Optional runtime leading dimension of matrix A. The lda and ldb must be specified together. If not specified, the compile-time lda is used.

  • ldb – Optional runtime leading dimension of matrix B. The lda and ldb must be specified together. If not specified, the compile-time ldb is used.