Skip to content

qttools.kernels.datastructure.cupy.dsdbcoo#

source module qttools.kernels.datastructure.cupy.dsdbcoo

Functions

source find_inds(self_rows: NDArray, self_cols: NDArray, rows: NDArray, cols: NDArray) → tuple[NDArray, NDArray, int]

Finds the corresponding indices of the given rows and columns.

Parameters

  • self_rows : NDArray — The rows of this matrix.

  • 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.

  • max_counts : int — The maximum number of matches found.

source 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.

source 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

source 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.

source 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.