warp.sparse.bsr_mm#

warp.sparse.bsr_mm(
x,
y,
z=None,
alpha=1.0,
beta=0.0,
masked=False,
work_arrays=None,
reuse_topology=False,
tile_size=0,
max_new_nnz=None,
)[source]#

Perform the sparse matrix-matrix multiplication z := alpha * x @ y + beta * z on BSR matrices x, y and z, and return z.

The x, y and z matrices are allowed to alias. If the matrix z is not provided as input, it will be allocated and treated as zero.

This method can be graph-captured if either:

  • masked=True

  • reuse_topology=True

  • max_new_nnz is provided

Parameters:

  • x (BsrMatrix[_MatrixBlockType[Rows, Any, Scalar] | _ScalarBlockType[Scalar]] | _BsrExpression[_MatrixBlockType[Rows, Any, Scalar] | _ScalarBlockType[Scalar]]) – Read-only left operand of the matrix-matrix product.

  • y (BsrMatrix[_MatrixBlockType[Any, Cols, Scalar] | _ScalarBlockType[Scalar]] | _BsrExpression[_MatrixBlockType[Any, Cols, Scalar] | _ScalarBlockType[Scalar]]) – Read-only right operand of the matrix-matrix product.

  • z (BsrMatrix[_MatrixBlockType[Rows, Cols, Scalar] | _ScalarBlockType[Scalar]] | None) – Mutable affine operand and result matrix. If z is not provided, it will be allocated and treated as zero.

  • alpha (Scalar) – Uniform scaling factor for the x @ y product

  • beta (Scalar) – Uniform scaling factor for z

  • masked (bool) – If True, keep the non-zero topology of z unchanged.

  • work_arrays (bsr_mm_work_arrays | None) – In most cases, this function will require the use of temporary storage. This storage can be reused across calls by passing an instance of bsr_mm_work_arrays in work_arrays.

  • reuse_topology (bool) – If True, reuse the product topology information stored in work_arrays rather than recompute it from scratch. The matrices x, y and z must be structurally similar to the previous call in which work_arrays were populated.

  • max_new_nnz (int | None) – If provided, the maximum number of non-zeros for the matrix-matrix product result (not counting the existing non-zeros in z).

  • tile_size (int) – If a positive integer, use tiles of this size to compute the matrix-matrix product. If negative, disable tile-based computation. Defaults to 0, which determines whether to use tiles using using an heuristic based on the matrix shape and number of non-zeros..

Return type:

BsrMatrix[_MatrixBlockType[Rows, Cols, Scalar] | _ScalarBlockType[Scalar]]