tilelang.utils.sparse¶

Attributes¶

GROUP_CONFIG

Functions¶

get_e_factor(a_dtype, meta_dtype)

Return how many a_dtype elements are indexed by one meta_dtype element.

get_e_replicate_factor(a_dtype)

Return how many consecutive threads share the same logical metadata value.

torch_compress(dense[, meta_dtype])

Reference 2:4 sparse compressor in pure PyTorch with natural row-major metadata. Modified from https://github.com/pytorch/pytorch/blob/bfa6895a345f6568624a4769238af6a9225e3fb8/torch/sparse/_semi_structured_conversions.py#L47

compress(A[, meta_dtype, block_m, block_k])

randn_semi_sparse(M, K[, dtype, device, transposed])

Generate a random semi-sparse tensor. The generated tensor will have 2:4 sparsity along the K dimension.

randint_semi_sparse(M, K, low, high[, dtype, device, ...])

Generate a random semi-sparse integer tensor. The generated tensor will have 2:4 sparsity along the K dimension.

arange_semi_sparse(M, K[, dtype, device, transposed])

Generate a semi-sparse tensor with values from 0 to M*K-1. The generated tensor will have 2:4 sparsity along the K dimension.

Module Contents¶

tilelang.utils.sparse.GROUP_CONFIG: dict[tilelang.language.dtypes.dtype, tuple[int, int]]¶
tilelang.utils.sparse.get_e_factor(a_dtype, meta_dtype)¶

Return how many a_dtype elements are indexed by one meta_dtype element.

Parameters:
Return type:

int

tilelang.utils.sparse.get_e_replicate_factor(a_dtype)¶

Return how many consecutive threads share the same logical metadata value.

Parameters:

a_dtype (tilelang.language.dtypes.dtype)

Return type:

int

tilelang.utils.sparse.torch_compress(dense, meta_dtype=None)¶

Reference 2:4 sparse compressor in pure PyTorch with natural row-major metadata. Modified from https://github.com/pytorch/pytorch/blob/bfa6895a345f6568624a4769238af6a9225e3fb8/torch/sparse/_semi_structured_conversions.py#L47

Each 4-bit chunk of the metadata integer encodes the two nonzero positions within one group of 4 consecutive elements:

bits [1:0] = index of first nonzero (0-3) bits [3:2] = index of second nonzero (0-3)

Parameters:

  • dense (torch.Tensor)

  • meta_dtype (torch.dtype | None)

Return type:

tuple[torch.Tensor, torch.Tensor]

tilelang.utils.sparse.compress(A, meta_dtype=None, block_m=None, block_k=None)¶
Parameters:

  • A (torch.Tensor)

  • meta_dtype (torch.dtype | None)

  • block_m (int | None)

  • block_k (int | None)

Return type:

tuple[torch.Tensor, torch.Tensor]

tilelang.utils.sparse.randn_semi_sparse(M, K, dtype=torch.float16, device='cuda', transposed=False)¶

Generate a random semi-sparse tensor. The generated tensor will have 2:4 sparsity along the K dimension. :param M: Number of rows :type M: int :param K: Number of columns :type K: int :param dtype: Data type of the tensor :param device: Device to create the tensor on :param transposed: If True, returns a transposed tensor of shape (K, M) :type transposed: bool

Parameters:

  • M (int)

  • K (int)

  • dtype (torch.dtype)

  • device (torch.device)

  • transposed (bool)

tilelang.utils.sparse.randint_semi_sparse(M, K, low, high, dtype=torch.int32, device='cuda', transposed=False)¶

Generate a random semi-sparse integer tensor. The generated tensor will have 2:4 sparsity along the K dimension. :param M: Number of rows :type M: int :param K: Number of columns :type K: int :param low: Lower bound of the random integers :type low: int :param high: Upper bound of the random integers :type high: int :param dtype: Data type of the tensor :param device: Device to create the tensor on :param transposed: If True, returns a transposed tensor of shape (K, M) :type transposed: bool

Parameters:

  • M (int)

  • K (int)

  • low (int)

  • high (int)

  • dtype (torch.dtype)

  • device (torch.device)

  • transposed (bool)

tilelang.utils.sparse.arange_semi_sparse(M, K, dtype=torch.float16, device='cuda', transposed=False)¶

Generate a semi-sparse tensor with values from 0 to M*K-1. The generated tensor will have 2:4 sparsity along the K dimension. :param M: Number of rows :type M: int :param K: Number of columns :type K: int :param dtype: Data type of the tensor :param device: Device to create the tensor on :param transposed: If True, returns a transposed tensor of shape (K, M) :type transposed: bool

Parameters:

  • M (int)

  • K (int)

  • dtype (torch.dtype)

  • device (torch.device)

  • transposed (bool)