warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/sparse.py

@@ -1,14 +1,29 @@
+from typing import Any, Generic, Optional, Tuple, TypeVar, Union
+
 import warp as wp
 import warp.types
 import warp.utils
+from warp.types import Array, Cols, Matrix, Rows, Scalar, Vector
+
+# typing hints
+
+_BlockType = TypeVar("BlockType")
+
+
+class _MatrixBlockType(Matrix):
+    pass
 
-from typing import Tuple, Any, Union
 
+class _ScalarBlockType(Generic[Scalar]):
+    pass
+
+
+BlockType = Union[_MatrixBlockType[Rows, Cols, Scalar], _ScalarBlockType[Scalar]]
 
 _struct_cache = dict()
 
 
-class BsrMatrix:
+class BsrMatrix(Generic[_BlockType]):
     """Untyped base class for BSR and CSR matrices.
 
     Should not be constructed directly but through functions such as :func:`bsr_zeros`.
@@ -16,15 +31,15 @@ class BsrMatrix:
     Attributes:
         nrow (int): Number of rows of blocks
         ncol (int): Number of columns of blocks
-        nnz (int):  Number of non-zero blocks: equal to `offsets[-1]`, cached on host for convenience
-        offsets (wp.array(dtype=int)): Array of size at least 1 + nrows containing start and end offsets og blocks in each row
-        columns (wp.array(dtype=int)): Array of size at least equal to nnz containing block column indices
-        values (wp.array(dtype=dtype)): Array of size at least equal to nnz containing block values
+        nnz (int):  Number of non-zero blocks: must be equal to ``offsets[nrow-1]``, cached on host for convenience
+        offsets (Array[int]): Array of size at least ``1 + nrows`` such that the start and end indices of the blocks of row ``r`` are ``offsets[r]`` and ``offsets[r+1]``, respectively.
+        columns (Array[int]): Array of size at least equal to ``nnz`` containing block column indices
+        values (Array[BlockType]): Array of size at least equal to ``nnz`` containing block values
     """
 
     @property
-    def scalar_type(self) -> type:
-        """Scalar type for each of the blocks' coefficients. FOr CSR matrices, this is equal to the block type"""
+    def scalar_type(self) -> Scalar:
+        """Scalar type for individual block coefficients. For CSR matrices, this is the same as the block type"""
         return warp.types.type_scalar_type(self.values.dtype)
 
     @property
@@ -33,20 +48,25 @@ class BsrMatrix:
         return getattr(self.values.dtype, "_shape_", (1, 1))
 
     @property
-    def block_size(self) -> Tuple[int, int]:
-        """Size of the individual blocks, i.e. number of rows per block times number of columsn per block"""
+    def block_size(self) -> int:
+        """Size of the individual blocks, i.e. number of rows per block times number of columns per block"""
         return warp.types.type_length(self.values.dtype)
 
     @property
     def shape(self) -> Tuple[int, int]:
-        """Shape of the matrix, i.e. number of rows/columns of blocks times number of rows/columsn per block"""
+        """Shape of the matrix, i.e. number of rows/columns of blocks times number of rows/columns per block"""
         block_shape = self.block_shape
         return (self.nrow * block_shape[0], self.ncol * block_shape[1])
 
 
-def bsr_matrix_t(dtype: type):
+def bsr_matrix_t(dtype: BlockType):
     dtype = wp.types.type_to_warp(dtype)
 
+    if not warp.types.type_is_matrix(dtype) and not dtype in warp.types.scalar_types:
+        raise ValueError(
+            f"BsrMatrix block type must be either warp matrix or scalar; got {warp.types.type_repr(dtype)}"
+        )
+
     class BsrMatrixTyped(BsrMatrix):
         nrow: int
         """Number of rows of blocks"""
@@ -79,11 +99,23 @@ def bsr_matrix_t(dtype: type):
 
 
 def bsr_zeros(
-    rows_of_blocks: int, cols_of_blocks: int, block_type: type, device: wp.context.Devicelike = None
+    rows_of_blocks: int,
+    cols_of_blocks: int,
+    block_type: BlockType,
+    device: wp.context.Devicelike = None,
 ) -> BsrMatrix:
     """
-    Constructs an empty BSR or CS matrix with the given shape
+    Constructs and returns an empty BSR or CSR matrix with the given shape
+
+    Args:
+        bsr: The BSR or CSR matrix to set to zero
+        rows_of_blocks: Number of rows of blocks
+        cols_of_blocks: Number of columns of blocks
+        block_type: Type of individual blocks. For CSR matrices, this should be a scalar type;
+                    for BSR matrices, this should be a matrix type (e.g. from :func:`warp.mat`)
+        device: Device on which to allocate the matrix arrays
     """
+
     bsr = bsr_matrix_t(block_type)()
 
     bsr.nrow = rows_of_blocks
@@ -110,19 +142,42 @@ def _bsr_ensure_fits(bsr: BsrMatrix, nrow: int = None, nnz: int = None):
         bsr.values = wp.empty(shape=(nnz,), dtype=bsr.values.dtype, device=bsr.values.device)
 
 
+def bsr_set_zero(bsr: BsrMatrix, rows_of_blocks: Optional[int] = None, cols_of_blocks: Optional[int] = None):
+    """
+    Sets a BSR matrix to zero, possibly changing its size
+
+    Args:
+        bsr: The BSR or CSR matrix to set to zero
+        rows_of_blocks: If not ``None``, the new number of rows of blocks
+        cols_of_blocks: If not ``None``, the new number of columns of blocks
+    """
+
+    if rows_of_blocks is not None:
+        bsr.nrow = rows_of_blocks
+    if cols_of_blocks is not None:
+        bsr.ncol = cols_of_blocks
+    bsr.nnz = 0
+    _bsr_ensure_fits(bsr)
+    bsr.offsets.zero_()
+
+
 def bsr_set_from_triplets(
-    dest: BsrMatrix,
-    rows: wp.array(dtype=int),
-    columns: wp.array(dtype=int),
-    values: wp.array(dtype=Any),
+    dest: BsrMatrix[BlockType[Rows, Cols, Scalar]],
+    rows: "Array[int]",
+    columns: "Array[int]",
+    values: "Array[Union[Scalar, BlockType[Rows, Cols, Scalar]]]",
 ):
     """
-    Fills a BSR matrix `dest` with values defined by COO triplets `rows`, `columns`, `values`.
+    Fills a BSR matrix with values defined by coordinate-oriented (COO) triplets, discarding existing blocks.
 
-    Values must be either one-dimensional with data type identical to the `dest` matrix block times,
-    or a 3d array with data type equal to the `dest` matrix scalar type.
+    The first dimension of the three input arrays must match, and determines the number of non-zeros in the constructed matrix.
 
-    Previous blocks of `dest` are discarded.
+    Args:
+        dest: Sparse matrix to populate
+        rows: Row index for each non-zero
+        columns: Columns index for each non-zero
+        values: Block values for each non-zero. Must be either a one-dimensional array with data type identical
+          to the `dest` matrix's block type, or a 3d array with data type equal to the `dest` matrix's scalar type.
     """
 
     if values.device != columns.device or values.device != rows.device or values.device != dest.values.device:
@@ -138,7 +193,7 @@ def bsr_set_from_triplets(
     elif values.ndim == 3:
         if values.shape[1:] != dest.block_shape:
             raise ValueError(
-                f"Last two dimensions in values array ({values.shape[1:]}) shoudl correspond to matrix block shape {(dest.block_shape)})"
+                f"Last two dimensions in values array ({values.shape[1:]}) should correspond to matrix block shape {(dest.block_shape)})"
             )
 
         if warp.types.type_scalar_type(values.dtype) != dest.scalar_type:
@@ -150,6 +205,9 @@ def bsr_set_from_triplets(
         raise ValueError("Number of dimension for values array should be 1 or 3")
 
     nnz = rows.shape[0]
+    if nnz == 0:
+        bsr_set_zero(dest)
+        return
 
     # Increase dest array sizes if needed
     _bsr_ensure_fits(dest, nnz=nnz)
@@ -186,8 +244,8 @@ def bsr_set_from_triplets(
     )
 
 
-def bsr_assign(dest: BsrMatrix, src: BsrMatrix):
-    """Copies the content of the `src` matrix to `dest`, possibly casting the block values."""
+def bsr_assign(dest: BsrMatrix[BlockType[Rows, Cols, Scalar]], src: BsrMatrix[BlockType[Rows, Cols, Any]]):
+    """Copies the content of the `src` matrix to `dest`, casting the block values if the two matrices use distinct scalar types."""
 
     if dest.values.device != src.values.device:
         raise ValueError("Source and destination matrices must reside on the same device")
@@ -207,8 +265,12 @@ def bsr_assign(dest: BsrMatrix, src: BsrMatrix):
         warp.utils.array_cast(out_array=dest.values, in_array=src.values, count=src.nnz)
 
 
-def bsr_copy(A: BsrMatrix, scalar_type=None):
-    """Returns a copy of matrix A, possibly casting values to a new scalar type"""
+def bsr_copy(A: BsrMatrix, scalar_type: Optional[Scalar] = None):
+    """Returns a copy of matrix ``A``, possibly changing its scalar type.
+
+    Args:
+       scalar_type: If provided, the returned matrix will use this scalar type instead of the one from `A`.
+    """
     if scalar_type is None:
         block_type = A.values.dtype
     elif A.block_shape == (1, 1):
@@ -221,7 +283,7 @@ def bsr_copy(A: BsrMatrix, scalar_type=None):
     return copy
 
 
-def bsr_set_transpose(dest: BsrMatrix, src: BsrMatrix):
+def bsr_set_transpose(dest: BsrMatrix[BlockType[Cols, Rows, Scalar]], src: BsrMatrix[BlockType[Rows, Cols, Scalar]]):
     """Assigns the transposed matrix `src` to matrix `dest`"""
 
     if dest.values.device != src.values.device:
@@ -230,10 +292,7 @@ def bsr_set_transpose(dest: BsrMatrix, src: BsrMatrix):
     if dest.scalar_type != src.scalar_type:
         raise ValueError("All arguments must have the same scalar type")
 
-    if src.block_shape == (1, 1):
-        transpose_block_shape = (1, 1)
-    else:
-        transpose_block_shape = src.block_shape[::-1]
+    transpose_block_shape = src.block_shape[::-1]
 
     if dest.block_shape != transpose_block_shape:
         raise ValueError(f"Destination block shape must be {transpose_block_shape}")
@@ -242,6 +301,9 @@ def bsr_set_transpose(dest: BsrMatrix, src: BsrMatrix):
     dest.ncol = src.nrow
     dest.nnz = src.nnz
 
+    if src.nnz == 0:
+        return
+
     # Increase dest array sizes if needed
     _bsr_ensure_fits(dest)
 
@@ -301,27 +363,33 @@ def _bsr_get_diag_kernel(
     end = A_offsets[row + 1]
 
     diag = wp.lower_bound(A_columns, beg, end, row)
-    if A_columns[diag] == row:
-        out[row] = A_values[diag]
+    if diag < end:
+        if A_columns[diag] == row:
+            out[row] = A_values[diag]
+
+
+def bsr_get_diag(A: BsrMatrix[_BlockType], out: "Optional[Array[BlockType]]" = None) -> "Array[BlockType]":
+    """Returns the array of blocks that constitute the diagonal of a sparse matrix.
 
+    Args:
+        A: the sparse matrix from which to extract the diagonal
+        out: if provided, the array into which to store the diagonal blocks
+    """
 
-def bsr_get_diag(A: BsrMatrix, out: wp.array = None):
-    """Returns the block diagonal of a square sparse matrix"""
-    if A.nrow != A.ncol:
-        raise ValueError("bsr_get_diag is only available for square sparse matrices")
+    dim = min(A.nrow, A.ncol)
 
     if out is None:
-        out = wp.zeros(shape=(A.nrow,), dtype=A.values.dtype, device=A.values.device)
+        out = wp.zeros(shape=(dim,), dtype=A.values.dtype, device=A.values.device)
     else:
         if out.dtype != A.values.dtype:
             raise ValueError(f"Output array must have type {A.values.dtype}")
         if out.device != A.values.device:
             raise ValueError(f"Output array must reside on device {A.values.device}")
-        if out.shape[0] < A.nrow:
-            raise ValueError(f"Output array must be of length at least {A.nrow}")
+        if out.shape[0] < dim:
+            raise ValueError(f"Output array must be of length at least {dim}")
 
     wp.launch(
-        kernel=_bsr_get_diag_kernel, dim=A.nrow, device=A.values.device, inputs=[A.offsets, A.columns, A.values, out]
+        kernel=_bsr_get_diag_kernel, dim=dim, device=A.values.device, inputs=[A.offsets, A.columns, A.values, out]
     )
 
     return out
@@ -329,40 +397,205 @@ def bsr_get_diag(A: BsrMatrix, out: wp.array = None):
 
 @wp.kernel
 def _bsr_set_diag_kernel(
+    diag: wp.array(dtype=Any),
+    A_offsets: wp.array(dtype=int),
+    A_columns: wp.array(dtype=int),
+    A_values: wp.array(dtype=Any),
+):
+    row = wp.tid()
+    A_offsets[row + 1] = row + 1
+    A_columns[row] = row
+    A_values[row] = diag[row]
+
+    if row == 0:
+        A_offsets[0] = 0
+
+
+@wp.kernel
+def _bsr_set_diag_constant_kernel(
+    diag_value: Any,
     A_offsets: wp.array(dtype=int),
     A_columns: wp.array(dtype=int),
+    A_values: wp.array(dtype=Any),
 ):
     row = wp.tid()
     A_offsets[row + 1] = row + 1
     A_columns[row] = row
+    A_values[row] = diag_value
 
     if row == 0:
         A_offsets[0] = 0
 
 
-def bsr_set_diag(A: BsrMatrix, diag: wp.array):
-    """Sets A as a block-diagonal square matrix"""
+def bsr_set_diag(
+    A: BsrMatrix[BlockType],
+    diag: "Union[BlockType, Array[BlockType]]",
+    rows_of_blocks: Optional[int] = None,
+    cols_of_blocks: Optional[int] = None,
+):
+    """Sets `A` as a block-diagonal matrix
+
+    Args:
+        A: the sparse matrix to modify
+        diag: Either a warp array of type ``A.values.dtype``, in which case each element will define one block of the diagonal,
+              or a constant value of type ``A.values.dtype``, in which case it will get assigned to all diagonal blocks.
+        rows_of_blocks: If not ``None``, the new number of rows of blocks
+        cols_of_blocks: If not ``None``, the new number of columns of blocks
+
+    The shape of the matrix will be defined one of the following, in that order:
+      - `rows_of_blocks` and `cols_of_blocks`, if provided. If only one is given, the second is assumed equal.
+      - the first dimension of `diag`, if `diag` is an array
+      - the current dimensions of `A` otherwise
+    """
+
+    if rows_of_blocks is None and cols_of_blocks is not None:
+        rows_of_blocks = cols_of_blocks
+    if cols_of_blocks is None and rows_of_blocks is not None:
+        cols_of_blocks = rows_of_blocks
+
+    if warp.types.is_array(diag):
+        if rows_of_blocks is None:
+            rows_of_blocks = diag.shape[0]
+            cols_of_blocks = diag.shape[0]
+
+    if rows_of_blocks is not None:
+        A.nrow = rows_of_blocks
+        A.ncol = cols_of_blocks
+
+    A.nnz = min(A.nrow, A.ncol)
+    _bsr_ensure_fits(A)
+
+    if warp.types.is_array(diag):
+        wp.launch(
+            kernel=_bsr_set_diag_kernel,
+            dim=A.nnz,
+            device=A.values.device,
+            inputs=[diag, A.offsets, A.columns, A.values],
+        )
+    else:
+        if not warp.types.type_is_value(type(diag)):
+            # Cast to launchable type
+            diag = A.values.dtype(diag)
+        wp.launch(
+            kernel=_bsr_set_diag_constant_kernel,
+            dim=A.nnz,
+            device=A.values.device,
+            inputs=[diag, A.offsets, A.columns, A.values],
+        )
 
-    A.nrow = diag.shape[0]
-    A.ncol = diag.shape[0]
-    A.nnz = diag.shape[0]
 
-    A.values = diag
-    if A.columns.size < A.nrow:
-        A.columns = wp.empty(shape=(A.nrow,), dtype=int, device=diag.device)
-    if A.offsets.size < A.nrow + 1:
-        A.offsets = wp.empty(shape=(A.nrow + 1,), dtype=int, device=diag.device)
+def bsr_diag(
+    diag: "Union[BlockType, Array[BlockType]]",
+    rows_of_blocks: Optional[int] = None,
+    cols_of_blocks: Optional[int] = None,
+) -> BsrMatrix["BlockType"]:
+    """Creates and returns a block-diagonal BSR matrix from an given block value or array of block values.
 
-    wp.launch(kernel=_bsr_set_diag_kernel, dim=A.nrow, device=A.values.device, inputs=[A.offsets, A.columns])
+    Args:
+        diag: Either a warp array of type ``A.values.dtype``, in which case each element will define one block of the diagonal,
+              or a constant value of type ``A.values.dtype``, in which case it will get assigned to all diagonal blocks.
+        rows_of_blocks: If not ``None``, the new number of rows of blocks
+        cols_of_blocks: If not ``None``, the new number of columns of blocks
 
+    The shape of the matrix will be defined one of the following, in that order:
+      - `rows_of_blocks` and `cols_of_blocks`, if provided. If only one is given, the second is assumed equal.
+      - the first dimension of `diag`, if `diag` is an array
+    """
+
+    if rows_of_blocks is None and cols_of_blocks is not None:
+        rows_of_blocks = cols_of_blocks
+    if cols_of_blocks is None and rows_of_blocks is not None:
+        cols_of_blocks = rows_of_blocks
+
+    if warp.types.is_array(diag):
+        if rows_of_blocks is None:
+            rows_of_blocks = diag.shape[0]
+            cols_of_blocks = diag.shape[0]
+
+        A = bsr_zeros(
+            rows_of_blocks,
+            cols_of_blocks,
+            block_type=diag.dtype,
+            device=diag.device,
+        )
+    else:
+        if rows_of_blocks is None:
+            raise ValueError(
+                "rows_of_blocks and/or cols_of_blocks must be provided for constructing a diagonal matrix with uniform diagonal"
+            )
+
+        block_type = type(diag)
+        if not warp.types.type_is_matrix(block_type) and len(getattr(diag, "shape", ())) == 2:
+            block_type = wp.mat(shape=diag.shape, dtype=diag.dtype)
+
+        A = bsr_zeros(
+            rows_of_blocks,
+            cols_of_blocks,
+            block_type=block_type,
+        )
 
-def bsr_diag(diag: wp.array):
-    """Creates a square block-diagonal BSR matrix from the values array `diag`"""
-    A = bsr_zeros(rows_of_blocks=diag.shape[0], cols_of_blocks=diag.shape[0], block_type=diag.dtype, device=diag.device)
     bsr_set_diag(A, diag)
     return A
 
 
+def bsr_set_identity(A: BsrMatrix, rows_of_blocks: Optional[int] = None):
+    """Sets `A` as the identity matrix
+
+    Args:
+        A: the sparse matrix to modify
+        rows_of_blocks: if provided, the matrix will be resized as a square matrix with `rows_of_blocks` rows and columns.
+    """
+
+    if A.block_shape == (1, 1):
+        identity = A.scalar_type(1.0)
+    else:
+        from numpy import eye
+
+        identity = eye(A.block_shape[0])
+
+    bsr_set_diag(A, diag=identity, rows_of_blocks=rows_of_blocks, cols_of_blocks=rows_of_blocks)
+
+
+def bsr_identity(
+    rows_of_blocks: int, block_type: BlockType[Rows, Rows, Scalar], device: wp.context.Devicelike = None
+) -> BsrMatrix[BlockType[Rows, Rows, Scalar]]:
+    """Creates and returns a square identity matrix.
+
+    Args:
+        rows_of_blocks: Number of rows and columns of blocks in the created matrix.
+        block_type: Block type for the newly created matrix -- must be square
+        device: Device onto which to allocate the data arrays
+    """
+    A = bsr_zeros(rows_of_blocks=rows_of_blocks, cols_of_blocks=rows_of_blocks, block_type=block_type, device=device)
+    bsr_set_identity(A)
+    return A
+
+
+@wp.kernel
+def _bsr_scale_kernel(
+    alpha: Any,
+    values: wp.array(dtype=Any),
+):
+    values[wp.tid()] = alpha * values[wp.tid()]
+
+
+def bsr_scale(x: BsrMatrix, alpha: Scalar) -> BsrMatrix:
+    """
+    Performs the operation ``x := alpha * x`` on BSR matrix `x` and returns `x`
+    """
+
+    if alpha != 1.0 and x.nnz > 0:
+        if alpha == 0.0:
+            bsr_set_zero(x)
+        else:
+            if not isinstance(alpha, x.scalar_type):
+                alpha = x.scalar_type(alpha)
+
+            wp.launch(kernel=_bsr_scale_kernel, dim=x.nnz, device=x.values.device, inputs=[alpha, x.values])
+
+    return x
+
+
 @wp.kernel
 def _bsr_get_block_row(dest_offset: int, bsr_offsets: wp.array(dtype=int), rows: wp.array(dtype=int)):
     i = wp.tid()
@@ -393,16 +626,75 @@ def _bsr_axpy_add_block(
     dst_values[block] = dst_values[block] + scale * src_values[i]
 
 
-def bsr_axpy(x: BsrMatrix, y: BsrMatrix, alpha: float = 1.0, beta: float = 1.0):
+class bsr_axpy_work_arrays:
+    """Opaque structure for persisting :func:`bsr_axpy` temporary work buffers across calls"""
+
+    def __init__(self):
+        self._reset(None)
+
+    def _reset(self, device):
+        self.device = device
+        self._sum_rows = None
+        self._sum_cols = None
+        self._old_y_values = None
+        self._old_x_values = None
+
+    def _allocate(self, device, y: BsrMatrix, sum_nnz: int):
+        if self.device != device:
+            self._reset(device)
+
+        if self._sum_rows is None or self._sum_rows.size < sum_nnz:
+            self._sum_rows = wp.empty(shape=(sum_nnz), dtype=int, device=self.device)
+        if self._sum_cols is None or self._sum_cols.size < sum_nnz:
+            self._sum_cols = wp.empty(shape=(sum_nnz), dtype=int, device=self.device)
+
+        if self._old_y_values is None or self._old_y_values.size < y.nnz:
+            self._old_y_values = wp.empty(shape=(y.nnz), dtype=y.values.dtype, device=self.device)
+
+
+def bsr_axpy(
+    x: BsrMatrix[BlockType[Rows, Cols, Scalar]],
+    y: Optional[BsrMatrix[BlockType[Rows, Cols, Scalar]]] = None,
+    alpha: Scalar = 1.0,
+    beta: Scalar = 1.0,
+    work_arrays: Optional[bsr_axpy_work_arrays] = None,
+) -> BsrMatrix[BlockType[Rows, Cols, Scalar]]:
     """
-    Performs the operation `y := alpha * X + beta * y` on BSR matrices `x` and `y`
+    Performs the sparse matrix addition ``y := alpha * X + beta * y`` on BSR matrices `x` and `y` and returns `y`.
+
+    The `x` and `y` matrices are allowed to alias.
+
+    Args:
+        x: Read-only right-hand-side.
+        y: Mutable left-hand-side. If `y` is not provided, it will be allocated and treated as zero.
+        alpha: Uniform scaling factor for `x`
+        beta: Uniform scaling factor for `y`
+        work_arrays: In most cases this function will require the use of temporary storage; this storage can be reused across calls by passing an instance of :class:`bsr_axpy_work_arrays` in `work_arrays`.
     """
 
     if y is None:
-        y = bsr_zeros(x.nrow, x.ncol, block_type=x.block_type, device=x.values.device)
+        # If not output matrix is provided, allocate it for convenience
+        y = bsr_zeros(x.nrow, x.ncol, block_type=x.values.dtype, device=x.values.device)
         beta = 0.0
 
-    device = y.values.device
+    # Handle easy cases first
+    if beta == 0.0 or y.nnz == 0:
+        bsr_assign(src=x, dest=y)
+        return bsr_scale(y, alpha=alpha)
+
+    if alpha == 0.0 or x.nnz == 0:
+        return bsr_scale(y, alpha=beta)
+
+    if not isinstance(alpha, y.scalar_type):
+        alpha = y.scalar_type(alpha)
+    if not isinstance(beta, y.scalar_type):
+        beta = y.scalar_type(beta)
+
+    if x == y:
+        # Aliasing case
+        return bsr_scale(y, alpha=alpha.value + beta.value)
+
+    # General case
 
     if x.values.device != y.values.device:
         raise ValueError("All arguments must reside on the same device")
@@ -413,20 +705,21 @@ def bsr_axpy(x: BsrMatrix, y: BsrMatrix, alpha: float = 1.0, beta: float = 1.0):
     if x.nrow != y.nrow or x.ncol != y.ncol:
         raise ValueError("Matrices must have the same number of rows and columns")
 
-    alpha = y.scalar_type(alpha)
-    beta = y.scalar_type(beta)
+    if work_arrays is None:
+        work_arrays = bsr_axpy_work_arrays()
 
     sum_nnz = x.nnz + y.nnz
-    sum_rows = wp.empty(shape=(sum_nnz), dtype=int, device=device)
-    sum_cols = wp.empty(shape=(sum_nnz), dtype=int, device=device)
+    device = y.values.device
+    work_arrays._allocate(device, y, sum_nnz)
+
+    wp.copy(work_arrays._sum_cols, y.columns, 0, 0, y.nnz)
+    wp.launch(kernel=_bsr_get_block_row, device=device, dim=y.nnz, inputs=[0, y.offsets, work_arrays._sum_rows])
 
-    if y.nnz > 0:
-        wp.copy(sum_cols, y.columns, 0, 0, y.nnz)
-        wp.launch(kernel=_bsr_get_block_row, device=device, dim=y.nnz, inputs=[0, y.offsets, sum_rows])
+    wp.copy(work_arrays._sum_cols, x.columns, y.nnz, 0, x.nnz)
+    wp.launch(kernel=_bsr_get_block_row, device=device, dim=x.nnz, inputs=[y.nnz, x.offsets, work_arrays._sum_rows])
 
-    if x.nnz > 0:
-        wp.copy(sum_cols, x.columns, y.nnz, 0, x.nnz)
-        wp.launch(kernel=_bsr_get_block_row, device=device, dim=x.nnz, inputs=[y.nnz, x.offsets, sum_rows])
+    # Save old y values before overwriting matrix
+    wp.copy(dest=work_arrays._old_y_values, src=y.values, count=y.nnz)
 
     # Increase dest array sizes if needed
     if y.columns.shape[0] < sum_nnz:
@@ -439,37 +732,55 @@ def bsr_axpy(x: BsrMatrix, y: BsrMatrix, alpha: float = 1.0, beta: float = 1.0):
     else:
         native_func = runtime.core.bsr_matrix_from_triplets_float_device
 
-    sum_nnz = native_func(
+    old_y_nnz = y.nnz
+    y.nnz = native_func(
         y.block_shape[0],
         y.block_shape[1],
         y.nrow,
         sum_nnz,
-        sum_rows.ptr,
-        sum_cols.ptr,
+        work_arrays._sum_rows.ptr,
+        work_arrays._sum_cols.ptr,
         0,
         y.offsets.ptr,
         y.columns.ptr,
         0,
     )
 
-    sum_values = wp.zeros(shape=(sum_nnz,), dtype=y.values.dtype, device=device)
+    _bsr_ensure_fits(y)
+    y.values.zero_()
 
     wp.launch(
         kernel=_bsr_axpy_add_block,
         device=device,
-        dim=y.nnz,
-        inputs=[0, beta, sum_rows, sum_cols, y.offsets, y.columns, y.values, sum_values],
+        dim=old_y_nnz,
+        inputs=[
+            0,
+            beta,
+            work_arrays._sum_rows,
+            work_arrays._sum_cols,
+            y.offsets,
+            y.columns,
+            work_arrays._old_y_values,
+            y.values,
+        ],
     )
+
     wp.launch(
         kernel=_bsr_axpy_add_block,
         device=device,
         dim=x.nnz,
-        inputs=[y.nnz, alpha, sum_rows, sum_cols, y.offsets, y.columns, x.values, sum_values],
+        inputs=[
+            old_y_nnz,
+            alpha,
+            work_arrays._sum_rows,
+            work_arrays._sum_cols,
+            y.offsets,
+            y.columns,
+            x.values,
+            y.values,
+        ],
     )
 
-    y.values = sum_values
-    y.nnz = sum_nnz
-
     return y
 
 
@@ -555,23 +866,77 @@ def _bsr_mm_compute_values(
             mm_values[mm_block] = mm_values[mm_block] + ax_val * y_values[y_block]
 
 
-_pinned_temp_count_buffer = {}
-
-
-def _get_pinned_temp_count_buffer(device):
-    device = str(device)
-    if device not in _pinned_temp_count_buffer:
-        _pinned_temp_count_buffer[device] = wp.empty(shape=(1,), dtype=int, pinned=True, device="cpu")
-
-    return _pinned_temp_count_buffer[device]
-
-
-def bsr_mm(x: BsrMatrix, y: BsrMatrix, z: BsrMatrix = None, alpha: float = 1.0, beta: float = 0.0):
+class bsr_mm_work_arrays:
+    """Opaque structure for persisting :func:`bsr_mm` temporary work buffers across calls"""
+
+    def __init__(self):
+        self._reset(None)
+
+    def _reset(self, device):
+        self.device = device
+        self._pinned_count_buffer = None
+        self._mm_row_counts = None
+        self._mm_rows = None
+        self._mm_cols = None
+        self._old_z_values = None
+        self._old_z_offsets = None
+        self._old_z_columns = None
+
+    def _allocate_stage_1(self, device, z: BsrMatrix, copied_z_nnz: int, z_aliasing: bool):
+        if self.device != device:
+            self._reset(device)
+
+        # Allocations that do not depend on any computation
+        if self.device.is_cuda:
+            if self._pinned_count_buffer is None:
+                self._pinned_count_buffer = wp.empty(shape=(1,), dtype=int, pinned=True, device="cpu")
+
+        if self._mm_row_counts is None or self._mm_row_counts.size < z.nrow + 1:
+            self._mm_row_counts = wp.empty(shape=(z.nrow + 1,), dtype=int, device=self.device)
+
+        if copied_z_nnz > 0:
+            if self._old_z_values is None or self._old_z_values.size < copied_z_nnz:
+                self._old_z_values = wp.empty(shape=(copied_z_nnz,), dtype=z.values.dtype, device=self.device)
+
+        if z_aliasing:
+            if self._old_z_columns is None or self._old_z_columns.size < z.nnz:
+                self._old_z_columns = wp.empty(shape=(z.nnz,), dtype=z.columns.dtype, device=self.device)
+            if self._old_z_offsets is None or self._old_z_offsets.size < z.nrow + 1:
+                self._old_z_offsets = wp.empty(shape=(z.nrow + 1,), dtype=z.offsets.dtype, device=self.device)
+
+    def _allocate_stage_2(self, mm_nnz: int):
+        # Allocations that depend on unmerged nnz estimate
+        if self._mm_rows is None or self._mm_rows.size < mm_nnz:
+            self._mm_rows = wp.empty(shape=(mm_nnz,), dtype=int, device=self.device)
+        if self._mm_cols is None or self._mm_cols.size < mm_nnz:
+            self._mm_cols = wp.empty(shape=(mm_nnz,), dtype=int, device=self.device)
+
+
+def bsr_mm(
+    x: BsrMatrix[BlockType[Rows, Any, Scalar]],
+    y: BsrMatrix[BlockType[Any, Cols, Scalar]],
+    z: Optional[BsrMatrix[BlockType[Rows, Cols, Scalar]]] = None,
+    alpha: Scalar = 1.0,
+    beta: Scalar = 0.0,
+    work_arrays: Optional[bsr_mm_work_arrays] = None,
+) -> BsrMatrix[BlockType[Rows, Cols, Scalar]]:
     """
-    Performs the operation `z := alpha * X * Y + beta * z` on BSR matrices `x`, `y` and `z`
+    Performs the sparse matrix-matrix multiplication ``z := alpha * x * y + beta * z`` on BSR matrices `x`, `y` and `z`, and returns `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.
+
+    Args:
+        x: Read-only left factor of the matrix-matrix product.
+        y: Read-only right factor of the matrix-matrix product.
+        z: Mutable left-hand-side. If `z` is not provided, it will be allocated and treated as zero.
+        alpha: Uniform scaling factor for the ``x * y`` product
+        beta: Uniform scaling factor for `z`
+        work_arrays: In most cases this function will require the use of temporary storage; this storage can be reused across calls by passing an instance of :class:`bsr_mm_work_arrays` in `work_arrays`.
     """
 
     if z is None:
+        # If not output matrix is provided, allocate it for convenience
         z_block_shape = (x.block_shape[0], y.block_shape[1])
         if z_block_shape == (1, 1):
             z_block_type = x.scalar_type
@@ -586,52 +951,85 @@ def bsr_mm(x: BsrMatrix, y: BsrMatrix, z: BsrMatrix = None, alpha: float = 1.0,
     if x.scalar_type != y.scalar_type or x.scalar_type != z.scalar_type:
         raise ValueError("Matrices must have the same scalar type")
 
-    if x.block_shape[0] != z.block_shape[0] or y.block_shape[1] != z.block_shape[1]:
-        raise ValueError("Incompatible blocks sizes for matrix multiplication")
+    if (
+        x.block_shape[0] != z.block_shape[0]
+        or y.block_shape[1] != z.block_shape[1]
+        or x.block_shape[1] != y.block_shape[0]
+    ):
+        raise ValueError("Incompatible block sizes for matrix multiplication")
 
-    if x.nrow != z.nrow or z.ncol != y.ncol:
+    if x.nrow != z.nrow or z.ncol != y.ncol or x.ncol != y.nrow:
         raise ValueError("Incompatible number of rows/columns for matrix multiplication")
 
     device = z.values.device
 
-    alpha = z.scalar_type(alpha)
-    beta = z.scalar_type(beta)
+    if alpha == 0.0 or x.nnz == 0 or y.nnz == 0:
+        # Easy case
+        return bsr_scale(z, beta)
+
+    if not isinstance(alpha, z.scalar_type):
+        alpha = z.scalar_type(alpha)
+    if not isinstance(beta, z.scalar_type):
+        beta = z.scalar_type(beta)
+
+    if work_arrays is None:
+        work_arrays = bsr_mm_work_arrays()
+
+    z_aliasing = z == x or z == y
+    copied_z_nnz = z.nnz if beta != 0.0 or z_aliasing else 0
+
+    work_arrays._allocate_stage_1(device, z, copied_z_nnz, z_aliasing)
 
     # Prefix sum of number of (unmerged) mm blocks per row
-    mm_row_counts = wp.empty(shape=(z.nrow + 1,), dtype=int, device=device)
     wp.launch(
         kernel=_bsr_mm_count_coeffs,
         device=device,
         dim=z.nrow,
-        inputs=[z.nnz, x.offsets, x.columns, y.offsets, mm_row_counts],
+        inputs=[copied_z_nnz, x.offsets, x.columns, y.offsets, work_arrays._mm_row_counts],
     )
-    warp.utils.array_scan(mm_row_counts, mm_row_counts)
+    warp.utils.array_scan(work_arrays._mm_row_counts, work_arrays._mm_row_counts)
 
     # Get back total counts on host
     if device.is_cuda:
-        mm_tot_count = _get_pinned_temp_count_buffer(device)
-        wp.copy(dest=mm_tot_count, src=mm_row_counts, src_offset=z.nrow, count=1)
-        wp.synchronize_stream(wp.get_stream())
-        mm_nnz = int(mm_tot_count.numpy()[0])
+        wp.copy(dest=work_arrays._pinned_count_buffer, src=work_arrays._mm_row_counts, src_offset=z.nrow, count=1)
+        wp.synchronize_stream(wp.get_stream(device))
+        mm_nnz = int(work_arrays._pinned_count_buffer.numpy()[0])
     else:
-        mm_nnz = int(mm_row_counts.numpy()[z.nrow])
+        mm_nnz = int(work_arrays._mm_row_counts.numpy()[z.nrow])
 
-    mm_rows = wp.empty(shape=(mm_nnz), dtype=int, device=device)
-    mm_cols = wp.empty(shape=(mm_nnz), dtype=int, device=device)
+    work_arrays._allocate_stage_2(mm_nnz)
 
-    # Copy z rows columns
-    wp.copy(mm_cols, z.columns, 0, 0, z.nnz)
-    wp.launch(kernel=_bsr_get_block_row, device=device, dim=z.nnz, inputs=[0, z.offsets, mm_rows])
+    # If z has a non-zero scale, save current data before overwriting it
+    if copied_z_nnz > 0:
+        # Copy z row and column indices
+        wp.copy(dest=work_arrays._mm_cols, src=z.columns, count=copied_z_nnz)
+        wp.launch(
+            kernel=_bsr_get_block_row, device=device, dim=copied_z_nnz, inputs=[0, z.offsets, work_arrays._mm_rows]
+        )
+        # Save current z values in temporary buffer
+        wp.copy(src=z.values, dest=work_arrays._old_z_values, count=copied_z_nnz)
+        if z_aliasing:
+            # If z is aliasing with x or y, need to save topology as well
+            wp.copy(src=z.columns, dest=work_arrays._old_z_columns, count=copied_z_nnz)
+            wp.copy(src=z.offsets, dest=work_arrays._old_z_offsets, count=z.nrow + 1)
 
     # Fill unmerged mm blocks rows and columns
     wp.launch(
         kernel=_bsr_mm_list_coeffs,
         device=device,
         dim=z.nrow,
-        inputs=[x.offsets, x.columns, y.offsets, y.columns, mm_row_counts, mm_rows, mm_cols],
+        inputs=[
+            x.offsets,
+            x.columns,
+            y.offsets,
+            y.columns,
+            work_arrays._mm_row_counts,
+            work_arrays._mm_rows,
+            work_arrays._mm_cols,
+        ],
     )
 
-    # Increase dest array sizes if needed
+    # Increase dest array size if needed
     if z.columns.shape[0] < mm_nnz:
         z.columns = wp.empty(shape=(mm_nnz,), dtype=int, device=device)
 
@@ -642,45 +1040,66 @@ def bsr_mm(x: BsrMatrix, y: BsrMatrix, z: BsrMatrix = None, alpha: float = 1.0,
     else:
         native_func = runtime.core.bsr_matrix_from_triplets_float_device
 
-    mm_nnz = native_func(
+    z.nnz = native_func(
         z.block_shape[0],
         z.block_shape[1],
         z.nrow,
         mm_nnz,
-        mm_rows.ptr,
-        mm_cols.ptr,
+        work_arrays._mm_rows.ptr,
+        work_arrays._mm_cols.ptr,
         0,
         z.offsets.ptr,
         z.columns.ptr,
         0,
     )
 
-    mm_values = wp.zeros(shape=(mm_nnz,), dtype=z.values.dtype, device=device)
-
-    # Copy blocks from z
-    wp.launch(
-        kernel=_bsr_axpy_add_block,
-        device=device,
-        dim=z.nnz,
-        inputs=[0, beta, mm_rows, mm_cols, z.offsets, z.columns, z.values, mm_values],
-    )
-
-    # Update z to point to result blocks
-    z.values = mm_values
-    z.nnz = mm_nnz
+    _bsr_ensure_fits(z)
+    z.values.zero_()
+
+    if copied_z_nnz > 0:
+        # Add back original z values
+        wp.launch(
+            kernel=_bsr_axpy_add_block,
+            device=device,
+            dim=copied_z_nnz,
+            inputs=[
+                0,
+                beta,
+                work_arrays._mm_rows,
+                work_arrays._mm_cols,
+                z.offsets,
+                z.columns,
+                work_arrays._old_z_values,
+                z.values,
+            ],
+        )
 
     # Add mm blocks to z values
-
-    if z.block_shape == (1, 1) and x.block_shape != (1, 1):
+    if (warp.types.type_is_matrix(x.values.dtype) or warp.types.type_is_matrix(y.values.dtype)) and not (
+        warp.types.type_is_matrix(z.values.dtype)
+    ):
         # Result block type is scalar, but operands are matrices
         # Cast result to (1x1) matrix to perform multiplication
-        mm_values = mm_values.view(wp.types.matrix(shape=(1, 1), dtype=z.scalar_type))
+        mm_values = z.values.view(wp.types.matrix(shape=(1, 1), dtype=z.scalar_type))
+    else:
+        mm_values = z.values
 
     wp.launch(
         kernel=_bsr_mm_compute_values,
         device=device,
         dim=z.nrow,
-        inputs=[alpha, x.offsets, x.columns, x.values, y.offsets, y.columns, y.values, z.offsets, z.columns, mm_values],
+        inputs=[
+            alpha,
+            work_arrays._old_z_offsets if x == z else x.offsets,
+            work_arrays._old_z_columns if x == z else x.columns,
+            work_arrays._old_z_values if x == z else x.values,
+            work_arrays._old_z_offsets if y == z else y.offsets,
+            work_arrays._old_z_columns if y == z else y.columns,
+            work_arrays._old_z_values if y == z else y.values,
+            z.offsets,
+            z.columns,
+            mm_values,
+        ],
     )
 
     return z
@@ -697,44 +1116,96 @@ def _bsr_mv_kernel(
     y: wp.array(dtype=Any),
 ):
     row = wp.tid()
-    beg = A_offsets[row]
-    end = A_offsets[row + 1]
 
-    yr = y[row]
-    v = yr - yr  # WAR to get zero with correct type
-    for block in range(beg, end):
-        v = v + A_values[block] * x[A_columns[block]]
+    # zero-initialize with type of y elements
+    scalar_zero = type(alpha)(0)
+    v = y.dtype(scalar_zero)
 
-    y[row] = beta * yr + alpha * v
+    if alpha != scalar_zero:
+        beg = A_offsets[row]
+        end = A_offsets[row + 1]
+        for block in range(beg, end):
+            v += A_values[block] * x[A_columns[block]]
+        v *= alpha
 
+    if beta != scalar_zero:
+        v += beta * y[row]
 
-def bsr_mv(A: BsrMatrix, x: wp.array, y: wp.array, alpha: float = 1.0, beta: float = 0.0):
+    y[row] = v
+
+
+def bsr_mv(
+    A: BsrMatrix[BlockType[Rows, Cols, Scalar]],
+    x: "Array[Vector[Cols, Scalar] | Scalar]",
+    y: Optional["Array[Vector[Rows, Scalar] | Scalar]"] = None,
+    alpha: Scalar = 1.0,
+    beta: Scalar = 0.0,
+    work_buffer: Optional["Array[Vector[Rows, Scalar] | Scalar]"] = None,
+) -> "Array[Vector[Rows, Scalar] | Scalar]":
     """
-    Naive implementation of sparse matrix-vector product, `y := alpha * A * x + beta * y`.
+    Performs the sparse matrix-vector product ``y := alpha * A * x + beta * y`` and returns `y`.
+
+    The `x` and `y` vectors are allowed to alias.
+
+    Args:
+        A: Read-only, left matrix factor of the matrix-vector product.
+        x: Read-only, right vector factor of the matrix-vector product.
+        y: Mutable left-hand-side. If `y` is not provided, it will be allocated and treated as zero.
+        alpha: Uniform scaling factor for `x`. If zero, `x` will not be read and may be left uninitialized.
+        beta: Uniform scaling factor for `y`. If zero, `y` will not be read and may be left uninitialized.
+        work_buffer: Temporary storage is required if and only if `x` and `y` are the same vector. If provided the `work_buffer` array
+            will be used for this purpose, otherwise a temporary allocation wil be performed.
     """
-    alpha = A.scalar_type(alpha)
-    beta = A.scalar_type(beta)
 
-    # if A.scalar_type != x.dtype or A.scalar_type != y.dtype:
-    #    raise ValueError("A, x and y must have the same data types")
+    if y is None:
+        # If no output array is provided, allocate one for convenience
+        y_vec_len = A.block_shape[0]
+        y_dtype = A.scalar_type if y_vec_len == 1 else wp.vec(length=y_vec_len, dtype=A.scalar_type)
+        y = wp.empty(shape=(A.nrow,), device=A.values.device, dtype=y_dtype)
+        y.zero_()
+        beta = 0.0
+
+    if not isinstance(alpha, A.scalar_type):
+        alpha = A.scalar_type(alpha)
+    if not isinstance(beta, A.scalar_type):
+        beta = A.scalar_type(beta)
 
     if A.values.device != x.device or A.values.device != y.device:
-        raise ValueError("A, x and y must reide on the same device")
+        raise ValueError("A, x and y must reside on the same device")
 
     if x.shape[0] != A.ncol:
         raise ValueError("Number of columns of A must match number of rows of x")
     if y.shape[0] != A.nrow:
         raise ValueError("Number of rows of A must match number of rows of y")
 
-    # Promote scalar vectors to length-1 vecs
-    block_shape = A.block_shape
-    if block_shape != (1, 1):
-        if block_shape[0] == 1:
+    if x == y:
+        # Aliasing case, need temporary storage
+        if work_buffer is None:
+            work_buffer = wp.empty_like(y)
+        elif work_buffer.size < y.size:
+            raise ValueError(f"Work buffer size is insufficient, needs to be at least {y.size}")
+        elif not wp.types.types_equal(work_buffer.dtype, y.dtype):
+            raise ValueError(f"Work buffer must have same data type as y, {wp.types.type_repr(y.dtype)}")
+
+        # Save old y values before overwriting vector
+        wp.copy(dest=work_buffer, src=y, count=y.size)
+        x = work_buffer
+
+    # Promote scalar vectors to length-1 vecs and conversely
+    if warp.types.type_is_matrix(A.values.dtype):
+        if A.block_shape[0] == 1:
             if y.dtype == A.scalar_type:
                 y = y.view(dtype=wp.vec(length=1, dtype=A.scalar_type))
-        if block_shape[1] == 1:
+        if A.block_shape[1] == 1:
             if x.dtype == A.scalar_type:
                 x = x.view(dtype=wp.vec(length=1, dtype=A.scalar_type))
+    else:
+        if A.block_shape[0] == 1:
+            if y.dtype != A.scalar_type:
+                y = y.view(dtype=A.scalar_type)
+        if A.block_shape[1] == 1:
+            if x.dtype != A.scalar_type:
+                x = x.view(dtype=A.scalar_type)
 
     wp.launch(
         kernel=_bsr_mv_kernel,
@@ -742,3 +1213,5 @@ def bsr_mv(A: BsrMatrix, x: wp.array, y: wp.array, alpha: float = 1.0, beta: flo
         dim=A.nrow,
         inputs=[alpha, A.offsets, A.columns, A.values, x, beta, y],
     )
+
+    return y