Skip to content

dsdbcsr#

Includes our CUDA csr datastructure kernels.

Functions:

  • find_inds –

    Finds the corresponding indices of the given rows and columns.

  • densify_block –

    Fills the dense block with the given data.

  • sparsify_block –

    Fills the data with the given dense block.

  • compute_rowptr_map –

    Computes the block-sorting index and the rowptr map.

find_inds #

find_inds(rowptr_map: dict, block_offsets: NDArray, self_cols: NDArray, rows: NDArray, cols: NDArray) -> tuple[NDArray, NDArray]

Finds the corresponding indices of the given rows and columns.

Parameters:

  • rowptr_map (dict) –

    The row pointer map.

  • block_offsets (NDArray) –

    The block offsets.

  • self_cols (NDArray) –

    The columns of this matrix.

  • rows (NDArray) –

    The rows to find the indices for.

  • cols (NDArray) –

    The columns to find the indices for.

Returns:

  • inds ( NDArray ) –

    The indices of the given rows and columns.

  • value_inds ( NDArray ) –

    The matching indices of this matrix.

densify_block #

densify_block(block: NDArray, block_offset: NDArray, self_cols: NDArray, rowptr: NDArray, data: NDArray)

Fills the dense block with the given data.

Parameters:

  • block (NDArray) –

    Preallocated dense block. Should be filled with zeros.

  • block_offset (NDArray) –

    The block offset.

  • self_cols (NDArray) –

    The column indices of this matrix.

  • rowptr (NDArray) –

    The row pointer of this matrix block.

  • data (NDArray) –

    The data to fill the block with.

sparsify_block #

sparsify_block(block: NDArray, block_offset: NDArray, self_cols: NDArray, rowptr: NDArray, data: NDArray)

Fills the data with the given dense block.

Parameters:

  • block (NDArray) –

    The dense block to sparsify.

  • block_offset (NDArray) –

    The block offset.

  • self_cols (NDArray) –

    The column indices of this matrix.

  • rowptr (NDArray) –

    The row pointer of this matrix block.

  • data (NDArray) –

    The data to be filled with the block.

compute_rowptr_map #

compute_rowptr_map(coo_rows: NDArray, coo_cols: NDArray, block_sizes: NDArray) -> dict

Computes the block-sorting index and the rowptr map.

Note

This is a combination of the bare block-sorting index computation and the rowptr map computation.

Parameters:

  • coo_rows (NDArray) –

    The row indices of the matrix in coordinate format.

  • coo_cols (NDArray) –

    The column indices of the matrix in coordinate format.

  • block_sizes (NDArray) –

    The block sizes of the block-sparse matrix we want to construct.

Returns:

  • sort_index ( NDArray ) –

    The block-sorting index for the sparse matrix.

  • rowptr_map ( dict ) –

    The row pointer map, describing the block-sparse matrix in blockwise column-sparse-row format.