warp.sparse#

Sparse linear algebra operations for Block Sparse Row (BSR) matrices.

This module provides GPU-accelerated sparse matrix operations for simulation, including matrix-matrix multiplication, matrix-vector multiplication, addition, scaling, and transpose operations. BSR format supports arbitrary block sizes, with Compressed Sparse Row (CSR) format supported as a special case using 1x1 blocks.

The BsrMatrix class supports operator overloading for intuitive matrix operations (+, -, *, @), and lower-level functions are available for fine-grained control over memory allocations.

Usage:

This module must be explicitly imported:

import warp.sparse

API#

`BsrMatrix`	Untyped base class for BSR and CSR matrices.
`bsr_axpy_work_arrays`	Opaque structure for persisting `bsr_axpy()` temporary work buffers across calls.
`bsr_mm_work_arrays`	Opaque structure for persisting `bsr_mm()` temporary work buffers across calls.
`bsr_assign`	Copy the content of the `src` BSR matrix to `dest`.
`bsr_axpy`	Perform the sparse matrix addition `y := alpha * X + beta * y` on BSR matrices `x` and `y` and return `y`.
`bsr_block_index`	Returns the index of the block at block-coordinates (row, col), or -1 if no such block exists.
`bsr_copy`	Return a copy of matrix `A`, possibly changing its scalar type.
`bsr_diag`	Create and return a block-diagonal BSR matrix from an given block value or array of block values.
`bsr_from_triplets`	Constructs a BSR matrix with values defined by coordinate-oriented (COO) triplets.
`bsr_get_diag`	Return the array of blocks that constitute the diagonal of a sparse matrix.
`bsr_identity`	Create and return a square identity matrix.
`bsr_matrix_t`
`bsr_mm`	Perform the sparse matrix-matrix multiplication `z := alpha * x @ y + beta * z` on BSR matrices `x`, `y` and `z`, and return `z`.
`bsr_mv`	Perform the sparse matrix-vector product `y := alpha * A * x + beta * y` and return `y`.
`bsr_row_index`	Returns the index of the row containing a given block, or -1 if no such row exists.
`bsr_scale`	Perform the operation `x := alpha * x` on BSR matrix `x` and return `x`.
`bsr_set_diag`	Set `A` as a block-diagonal matrix.
`bsr_set_from_triplets`	Fill a BSR matrix with values defined by coordinate-oriented (COO) triplets, discarding existing blocks.
`bsr_set_identity`	Set `A` as the identity matrix.
`bsr_set_transpose`	Assign the transposed matrix `src` to matrix `dest`.
`bsr_set_zero`	Set a BSR matrix to zero, possibly changing its size.
`bsr_transposed`	Return a copy of the transposed matrix `A`.
`bsr_zeros`	Construct and return an empty BSR or CSR matrix with the given shape.