Skip to content

dsdbcoo#

Includes our CUDA coo datastructure kernels.

Functions:

compute_block_slice #

compute_block_slice(rows: NDArray, cols: NDArray, block_offsets: NDArray, row: int, col: int) -> slice

Computes the slice of the block in the data.

Parameters:

  • rows (NDArray) –

    The row indices of the matrix.

  • cols (NDArray) –

    The column indices of the matrix.

  • block_offsets (NDArray) –

    The offsets of the blocks.

  • row (int) –

    The block row to compute the slice for.

  • col (int) –

    The block column to compute the slice for.

Returns:

  • start ( int ) –

    The start index of the block.

  • stop ( int ) –

    The stop index of the block.

densify_block #

densify_block(block: NDArray, rows: NDArray, cols: NDArray, data: NDArray, block_slice: slice, row_offset: int, col_offset: int, use_kernel: bool = QTX_USE_DENSIFY_BLOCK)

Fills the dense block with the given data.

Note

This is not a raw kernel, as there seems to be no performance gain for this operation on the GPU.

Parameters:

  • rows (NDArray) –

    The rows at which to fill the block.

  • cols (NDArray) –

    The columns at which to fill the block.

  • data (NDArray) –

    The data to fill the block with.

  • block (NDArray) –

    Preallocated dense block. Should be filled with zeros.

  • block_slice (slice) –

    The slice of the block to fill.

  • row_offset (int) –

    The row offset of the block.

  • col_offset (int) –

    The column offset of the block

sparsify_block #

sparsify_block(block: NDArray, rows: NDArray, cols: NDArray, data: NDArray)

Fills the data with the given dense block.

Note

This is not a raw kernel, as there seems to be no performance gain for this operation on the GPU.

Parameters:

  • block (NDArray) –

    The dense block to sparsify.

  • rows (NDArray) –

    The rows at which to fill the block.

  • cols (NDArray) –

    The columns at which to fill the block.

  • data (NDArray) –

    The data to be filled with the block.

compute_block_sort_index #

compute_block_sort_index(coo_rows: NDArray, coo_cols: NDArray, block_sizes: NDArray) -> NDArray

Computes the block-sorting index for a sparse matrix.

Note

Due to the Python for loop around the kernel, this method will perform best for larger block sizes (>500).

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 indexing that sorts the data by block-row and -column.