warp.sparse#

Warp includes a sparse linear algebra module warp.sparse that implements common sparse matrix operations for simulation. The module provides efficient implementations of Block Sparse Row (BSR) matrices, with Compressed Sparse Row (CSR) matrices supported as a special case with 1x1 block size.

Working with Sparse Matrices#

Creating Sparse Matrices#

The most common way to create a sparse matrix is using coordinate (COO) format triplets:

import warp as wp
from warp.sparse import bsr_from_triplets

# Create a 4x4 sparse matrix with 2x2 blocks
rows = wp.array([0, 1, 2], dtype=int)  # Row indices
cols = wp.array([1, 2, 3], dtype=int)  # Column indices
vals = wp.array([...], dtype=float)    # Block values

# Create BSR matrix
A = bsr_from_triplets(
    rows_of_blocks=2,      # Number of rows of blocks
    cols_of_blocks=2,      # Number of columns of blocks
    rows=rows,            # Row indices
    columns=cols,         # Column indices
    values=vals           # Block values
)

You can also create special matrices:

from warp.sparse import bsr_zeros, bsr_identity, bsr_diag

# Create empty matrix
A = bsr_zeros(rows_of_blocks=4, cols_of_blocks=4, block_type=wp.float32)

# Create identity matrix
I = bsr_identity(rows_of_blocks=4, block_type=wp.float32)

# Create diagonal matrix
D = bsr_diag(diag=wp.array([1.0, 2.0, 3.0, 4.0]))

Block Sizes#

BSR matrices support different block sizes. For example:

# 1x1 blocks (CSR format)
A = bsr_zeros(4, 4, block_type=wp.float32)

# 3x3 blocks
A = bsr_zeros(2, 2, block_type=wp.mat33)

# Rectangular block sizes
A = bsr_zeros(2, 3, block_type=wp.mat23)

The module also provides functions to convert between different block sizes and scalar types:

from warp.sparse import bsr_copy, bsr_assign

# Convert block size from 2x2 to 1x1 (BSR to CSR)
A_csr = bsr_copy(A, block_shape=(1, 1))

# Convert block size from 1x1 to 2x2 (CSR to BSR)
A_bsr = bsr_copy(A, block_shape=(2, 2))

# Convert scalar type from float32 to float64
A_double = bsr_copy(A, scalar_type=wp.float64)

# Convert both block size and scalar type
A_new = bsr_copy(A, block_shape=(2, 2), scalar_type=wp.float64)

# In-place conversion using bsr_assign
B = bsr_zeros(rows_of_blocks=4, cols_of_blocks=4, block_type=wp.mat22)
bsr_assign(src=A, dest=B)  # Converts A to 2x2 blocks and stores in B

Note

When converting block sizes, the source and destination block dimensions must be compatible:

The new block dimensions must either be multiples or exact dividers of the original dimensions
The total matrix dimensions must be evenly divisible by the new block size

For example:

# Valid conversions:
A = bsr_zeros(4, 4, block_type=wp.mat22)  # 8x8 matrix with 2x2 blocks
B = bsr_copy(A, block_shape=(1, 1))       # 8x8 matrix with 1x1 blocks
C = bsr_copy(A, block_shape=(4, 4))       # 8x8 matrix with 4x4 blocks

# Invalid conversion (will raise ValueError):
D = bsr_copy(A, block_shape=(3, 3))       # 3x3 blocks don't divide 8x8 evenly

Non-Zero Block Count#

The number of non-zero blocks in a BSR matrix is computed on the device and not automatically synchronized to the host to avoid performance overhead and allow graph capture. The nnz attribute of a BSR matrix is always an upper bound for the number of non-zero blocks, but the actual count is stored on the device at offsets[nrow].

To get the exact count on host, you can explicitly synchronize using the nnz_sync() method:

# The nnz attribute is an upper bound
upper_bound = A.nnz

# Get the exact number of non-zero blocks
exact_nnz = A.nnz_sync()

Note

The nnz_sync() method ensures that any ongoing transfer of the exact nnz number from the device offsets array to the host has completed and updates the nnz upper bound. This synchronization is only necessary when you need the exact count - for most operations, the upper bound is sufficient.

If the number of non-zeros has been changed from outside of the warp.sparse builtin functions, for instance by direct modifications to the offsets array, use the notify_nnz_changed() method to ensure consistency.

Converting back to COO Format#

You can convert a BSR matrix back to coordinate (COO) format using the matrix’s properties:

# Get row indices from compressed format
nnz = A.nnz_sync()
rows = A.uncompress_rows()[:nnz]  # Returns array of row indices for each block

# Get column indices and values
cols = A.columns[:nnz]            # Column indices for each block
vals = A.values[:nnz]             # Block values

# Now you have the COO format:
# - rows[i] is the row index of block i
# - cols[i] is the column index of block i
# - vals[i] is the value of block i

Matrix Operations#

The module supports common matrix operations with overloaded operators:

# Matrix addition
C = A + B

# Matrix subtraction
C = A - B

# Scalar multiplication
C = 2.0 * A

# Matrix multiplication
C = A @ B

# Matrix-vector multiplication
y = A @ x

# Transpose
AT = A.transpose()

# In-place operations
A += B
A -= B
A *= 2.0
A @= B

For more control about memory allocations, you may use the underlying lower-level functions:

from warp.sparse import bsr_mm, bsr_mv, bsr_axpy, bsr_scale

# Matrix-matrix multiplication
# C = alpha * A @ B + beta * C
bsr_mm(
    A, B, C,           # Input and output matrices
    alpha=1.0,         # Scale factor for A @ B
    beta=0.0,          # Scale factor for C
)

# Matrix-vector multiplication
# y = alpha * A @ x + beta * y
bsr_mv(
    A, x, y,           # Input matrix, vector, and output vector
    alpha=1.0,         # Scale factor for A @ x
    beta=0.0,          # Scale factor for y
)

# Matrix addition
# C = alpha * A + beta * B
bsr_axpy(
    A, B, C,           # Input and output matrices
    alpha=1.0,         # Scale factor for A
    beta=1.0,          # Scale factor for B
)

# Matrix scaling
# A = alpha * A
bsr_scale(
    A,                 # Input/output matrix
    alpha,             # Scale factor
)

Kernel-level utilities#

The following Warp functions are available for use in kernels:

warp.sparse.bsr_row_index(offsets, row_count, block_index)#

Returns the index of the row containing a given block, or -1 if no such row exists.

Parameters:

offsets (array(ndim=1, dtype=int32)) – Array of size at least 1 + row_count containing the offsets of the blocks in each row.
row_count (int) – Number of rows of blocks.
block_index (int) – Index of the block.

Return type:

int

warp.sparse.bsr_block_index(row, col, bsr_offsets, bsr_columns)#

Returns the index of the block at block-coordinates (row, col), or -1 if no such block exists. Assumes that the segments of bsr_columns corresponding to each row are sorted.

Parameters:

row (int) – Row of the block.
col (int) – Column of the block.
bsr_offsets (array(ndim=1, dtype=int32)) – Array of size at least 1 + row containing the offsets of the blocks in each row.
bsr_columns (array(ndim=1, dtype=int32)) – Array of size at least equal to bsr_offsets[row + 1] containing the column indices of the blocks.

Return type:

int

Reference#

class warp.sparse.BsrMatrix[source]#

Untyped base class for BSR and CSR matrices.

Should not be constructed directly but through functions such as bsr_zeros().

nrow#

Number of rows of blocks.

Type:: int

ncol#

Number of columns of blocks.

Type:: int

nnz#

Upper bound for the number of non-zero blocks, used for dimensioning launches. The exact number is at offsets[nrow-1]. See also nnz_sync().

Type:: int

offsets#

Array of size at least 1 + nrow such that the start and end indices of the blocks of row r are offsets[r] and offsets[r+1], respectively.

Type:: Array[int]

columns#

Array of size at least equal to nnz containing block column indices.

Type:: Array[int]

values#

Array of size at least equal to nnz containing block values.

Type:: Array[BlockType]

property block_shape: tuple[int, int][source]#: Shape of the individual blocks.

property block_size: int[source]#: Size of the individual blocks, i.e. number of rows per block times number of columns per block.

copy_nnz_async()[source]#

Starts the asynchronous transfer of the exact nnz from the device offsets array to host and records an event for completion.

Deprecated; prefer notify_nnz_changed() instead, which will make sure to resize arrays if necessary.

Return type:: None

property device: Device[source]#: Device on which offsets, columns, and values are allocated – assumed to be the same for all three arrays.

property dtype: type[source]#: Data type for individual block values.

nnz_sync()[source]#

Synchronize the number of non-zeros from the device offsets array to the host.

Ensures that any ongoing transfer of the exact nnz number from the device offsets array to the host has completed, or, if none has been scheduled yet, starts a new transfer and waits for it to complete.

Then updates the host-side nnz upper bound to match the exact one, and returns it.

Iterative Linear Solvers#

Warp provides several iterative linear solvers with optional preconditioning:

Conjugate Gradient (CG)
Conjugate Residual (CR)
Bi-Conjugate Gradient Stabilized (BiCGSTAB)
Generalized Minimal Residual (GMRES)

from warp.optim.linear import cg

# Solve Ax = b
x = cg(A, b, max_iter=100, tol=1e-6)

Note

While primarily intended for sparse matrices, these solvers also work with dense linear operators provided as 2D Warp arrays. Custom operators can be implemented using the LinearOperator interface.

class warp.optim.linear.LinearOperator(shape, dtype, device, matvec)[source]#

Linear operator to be used as left-hand-side of linear iterative solvers.

Parameters:

shape (Tuple[int, int]) – Tuple containing the number of rows and columns of the operator
dtype (type) – Type of the operator elements
device (Device) – Device on which computations involving the operator should be performed
matvec (Callable) – Matrix-vector multiplication routine

The matrix-vector multiplication routine should have the following signature:

def matvec(x: wp.array, y: wp.array, z: wp.array, alpha: Scalar, beta: Scalar):
    '''Perform a generalized matrix-vector product.

    This function computes the operation z = alpha * (A @ x) + beta * y, where 'A'
    is the linear operator represented by this class.
    '''
    ...

For performance reasons, by default the iterative linear solvers in this module will try to capture the calls for one or more iterations in CUDA graphs. If the matvec routine of a custom LinearOperator cannot be graph-captured, the use_cuda_graph=False parameter should be passed to the solver function.

__init__(shape, dtype, device, matvec)[source]#

Parameters:

shape (Tuple[int, int])
dtype (type)
device (Device)
matvec (Callable)

property device: Device[source]#

property dtype: type[source]#

property matvec: Callable[source]#

property scalar_type[source]#

property shape: Tuple[int, int][source]#

warp.optim.linear.aslinearoperator(A)[source]#

Casts the dense or sparse matrix A as a LinearOperator

A must be of one of the following types:

warp.sparse.BsrMatrix

two-dimensional warp.array; then A is assumed to be a dense matrix

one-dimensional warp.array; then A is assumed to be a diagonal matrix

warp.sparse.LinearOperator; no casting necessary

Parameters:: A (array | BsrMatrix | LinearOperator)
Return type:: LinearOperator

warp.optim.linear.bicgstab( A, b, x, tol=None, atol=None, maxiter=0, M=None, callback=None, check_every=10, use_cuda_graph=True, is_left_preconditioner=False, )[source]#

Computes an approximate solution to a linear system using the Biconjugate Gradient Stabilized method (BiCGSTAB).

Parameters:

A (array | BsrMatrix | LinearOperator) – the linear system’s left-hand-side
b (array) – the linear system’s right-hand-side
x (array) – initial guess and solution vector
tol (float | None) – relative tolerance for the residual, as a ratio of the right-hand-side norm
atol (float | None) – absolute tolerance for the residual
maxiter (float | None) – maximum number of iterations to perform before aborting. Defaults to the system size.
M (array | BsrMatrix | LinearOperator | None) – optional left- or right-preconditioner, ideally chosen such that M A (resp A M) is close to identity.
callback (Callable | None) – function to be called every check_every iteration with the current iteration number, residual and tolerance. If check_every is 0, the callback should be a Warp kernel.
check_every – number of iterations every which to call callback, check the residual against the tolerance and possibility terminate the algorithm. Setting check_every to 0 disables host-side residual checks, making the solver fully CUDA-graph capturable. If conditional CUDA graphs are supported, convergence checks are performed device-side; otherwise, the solver will always run to the maximum number of iterations.
use_cuda_graph – If true and when run on a CUDA device, capture the solver iteration as a CUDA graph for reduced launch overhead. The linear operator and preconditioner must only perform graph-friendly operations.
is_left_preconditioner – whether M should be used as a left- or right- preconditioner.

Returns:

Tuple (final_iteration, residual_norm, absolute_tolerance)

final_iteration: The number of iterations performed before convergence or reaching maxiter
residual_norm: The final residual norm ||b - Ax||
absolute_tolerance: The absolute tolerance used for convergence checking

If check_every is 0: Tuple (final_iteration_array, residual_norm_squared_array, absolute_tolerance_squared_array)

final_iteration_array: Device array containing the number of iterations performed
residual_norm_squared_array: Device array containing the squared residual norm ||b - Ax||²
absolute_tolerance_squared_array: Device array containing the squared absolute tolerance

Return type:

If check_every > 0

If both tol and atol are provided, the absolute tolerance used as the termination criterion for the residual norm is max(atol, tol * norm(b)).

warp.optim.linear.cg( A, b, x, tol=None, atol=None, maxiter=0, M=None, callback=None, check_every=10, use_cuda_graph=True, )[source]#

Computes an approximate solution to a symmetric, positive-definite linear system using the Conjugate Gradient algorithm.

Parameters:

A (array | BsrMatrix | LinearOperator) – the linear system’s left-hand-side
b (array) – the linear system’s right-hand-side
x (array) – initial guess and solution vector
tol (float | None) – relative tolerance for the residual, as a ratio of the right-hand-side norm
atol (float | None) – absolute tolerance for the residual
maxiter (float | None) – maximum number of iterations to perform before aborting. Defaults to the system size.
M (array | BsrMatrix | LinearOperator | None) – optional left-preconditioner, ideally chosen such that M A is close to identity.
callback (Callable | None) – function to be called every check_every iteration with the current iteration number, residual and tolerance. If check_every is 0, the callback should be a Warp kernel.
check_every – number of iterations every which to call callback, check the residual against the tolerance and possibility terminate the algorithm. Setting check_every to 0 disables host-side residual checks, making the solver fully CUDA-graph capturable. If conditional CUDA graphs are supported, convergence checks are performed device-side; otherwise, the solver will always run to the maximum number of iterations.
use_cuda_graph – If true and when run on a CUDA device, capture the solver iteration as a CUDA graph for reduced launch overhead. The linear operator and preconditioner must only perform graph-friendly operations.

Returns:

Tuple (final_iteration, residual_norm, absolute_tolerance)

final_iteration: The number of iterations performed before convergence or reaching maxiter
residual_norm: The final residual norm ||b - Ax||
absolute_tolerance: The absolute tolerance used for convergence checking

If check_every is 0: Tuple (final_iteration_array, residual_norm_squared_array, absolute_tolerance_squared_array)

final_iteration_array: Device array containing the number of iterations performed
residual_norm_squared_array: Device array containing the squared residual norm ||b - Ax||²
absolute_tolerance_squared_array: Device array containing the squared absolute tolerance

Return type:

If check_every > 0

If both tol and atol are provided, the absolute tolerance used as the termination criterion for the residual norm is max(atol, tol * norm(b)).

warp.optim.linear.cr( A, b, x, tol=None, atol=None, maxiter=0, M=None, callback=None, check_every=10, use_cuda_graph=True, )[source]#

Computes an approximate solution to a symmetric, positive-definite linear system using the Conjugate Residual algorithm.

Parameters:

A (array | BsrMatrix | LinearOperator) – the linear system’s left-hand-side
b (array) – the linear system’s right-hand-side
x (array) – initial guess and solution vector
tol (float | None) – relative tolerance for the residual, as a ratio of the right-hand-side norm
atol (float | None) – absolute tolerance for the residual
maxiter (float | None) – maximum number of iterations to perform before aborting. Defaults to the system size. Note that the current implementation always performs iterations in pairs, and as a result may exceed the specified maximum number of iterations by one.
M (array | BsrMatrix | LinearOperator | None) – optional left-preconditioner, ideally chosen such that M A is close to identity.
callback (Callable | None) – function to be called every check_every iteration with the current iteration number, residual and tolerance. If check_every is 0, the callback should be a Warp kernel.
check_every – number of iterations every which to call callback, check the residual against the tolerance and possibility terminate the algorithm. Setting check_every to 0 disables host-side residual checks, making the solver fully CUDA-graph capturable. If conditional CUDA graphs are supported, convergence checks are performed device-side; otherwise, the solver will always run to the maximum number of iterations.
use_cuda_graph – If true and when run on a CUDA device, capture the solver iteration as a CUDA graph for reduced launch overhead. The linear operator and preconditioner must only perform graph-friendly operations.

Returns:

Tuple (final_iteration, residual_norm, absolute_tolerance)

final_iteration: The number of iterations performed before convergence or reaching maxiter
residual_norm: The final residual norm ||b - Ax||
absolute_tolerance: The absolute tolerance used for convergence checking

If check_every is 0: Tuple (final_iteration_array, residual_norm_squared_array, absolute_tolerance_squared_array)

final_iteration_array: Device array containing the number of iterations performed
residual_norm_squared_array: Device array containing the squared residual norm ||b - Ax||²
absolute_tolerance_squared_array: Device array containing the squared absolute tolerance

Return type:

If check_every > 0

If both tol and atol are provided, the absolute tolerance used as the termination criterion for the residual norm is max(atol, tol * norm(b)).

warp.optim.linear.gmres( A, b, x, tol=None, atol=None, restart=31, maxiter=0, M=None, callback=None, check_every=31, use_cuda_graph=True, is_left_preconditioner=False, )[source]#

Computes an approximate solution to a linear system using the restarted Generalized Minimum Residual method (GMRES[k]).

Parameters:

A (array | BsrMatrix | LinearOperator) – the linear system’s left-hand-side
b (array) – the linear system’s right-hand-side
x (array) – initial guess and solution vector
tol (float | None) – relative tolerance for the residual, as a ratio of the right-hand-side norm
atol (float | None) – absolute tolerance for the residual
restart – The restart parameter, i.e, the k in GMRES[k]. In general, increasing this parameter reduces the number of iterations but increases memory consumption.
maxiter (float | None) – maximum number of iterations to perform before aborting. Defaults to the system size. Note that the current implementation always perform restart iterations at a time, and as a result may exceed the specified maximum number of iterations by restart-1.
M (array | BsrMatrix | LinearOperator | None) – optional left- or right-preconditioner, ideally chosen such that M A (resp A M) is close to identity.
callback (Callable | None) – function to be called every check_every iteration with the current iteration number, residual and tolerance. If check_every is 0, the callback should be a Warp kernel.
check_every – number of iterations every which to call callback, check the residual against the tolerance and possibility terminate the algorithm. Setting check_every to 0 disables host-side residual checks, making the solver fully CUDA-graph capturable. If conditional CUDA graphs are supported, convergence checks are performed device-side; otherwise, the solver will always run to the maximum number of iterations.
use_cuda_graph – If true and when run on a CUDA device, capture the solver iteration as a CUDA graph for reduced launch overhead. The linear operator and preconditioner must only perform graph-friendly operations.
is_left_preconditioner – whether M should be used as a left- or right- preconditioner.

Returns:

Tuple (final_iteration, residual_norm, absolute_tolerance)

final_iteration: The number of iterations performed before convergence or reaching maxiter
residual_norm: The final residual norm ||b - Ax||
absolute_tolerance: The absolute tolerance used for convergence checking

If check_every is 0: Tuple (final_iteration_array, residual_norm_squared_array, absolute_tolerance_squared_array)

final_iteration_array: Device array containing the number of iterations performed
residual_norm_squared_array: Device array containing the squared residual norm ||b - Ax||²
absolute_tolerance_squared_array: Device array containing the squared absolute tolerance

Return type:

If check_every > 0

If both tol and atol are provided, the absolute tolerance used as the termination criterion for the residual norm is max(atol, tol * norm(b)).

warp.optim.linear.preconditioner(A, ptype='diag')[source]#

Constructs and returns a preconditioner for an input matrix.

Parameters:

A (array | BsrMatrix | LinearOperator) – The matrix for which to build the preconditioner
ptype (str) –
The type of preconditioner. Currently the following values are supported:
- "diag": Diagonal (a.k.a. Jacobi) preconditioner
- "diag_abs": Similar to Jacobi, but using the absolute value of diagonal coefficients
- "id": Identity (null) preconditioner

Return type:

LinearOperator