tilelang.utils.sparse
=====================

.. py:module:: tilelang.utils.sparse


Attributes
----------

.. autoapisummary::

   tilelang.utils.sparse.compress_util


Functions
---------

.. autoapisummary::

   tilelang.utils.sparse.compress_sm90
   tilelang.utils.sparse.compress_sm80
   tilelang.utils.sparse.compress
   tilelang.utils.sparse.randn_semi_sparse
   tilelang.utils.sparse.randint_semi_sparse
   tilelang.utils.sparse.arange_semi_sparse


Module Contents
---------------

.. py:data:: compress_util

.. py:function:: compress_sm90(A, block_k, transposed)

.. py:function:: compress_sm80(A, transposed)

.. py:function:: compress(A, transposed, arch = None, **kwargs)

   Compress a tensor using the appropriate method based on the CUDA architecture.


.. py:function:: 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


.. py:function:: 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


.. py:function:: 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