Update on "[sparse] Add fast semi-structured spasification kernels"

jcaip · jcaip · commit b3898ad6e513 · 2024-04-18T11:48:41.000-07:00
This PR adds in fast semi-structured sparsification kernels to PyTorch. These kernels allow for accelerated semi-structured sparsification kernels in PyTorch. The kernels have been added as aten native functions In particular, three new functions have been added: * `torch._sparse_semi_structured_tile` This function will return the packed representation and metadata for both X and X', as well as the thread masks. Note that this applies 2:4 sparsity in a 4x4 tile instead of a 1x4 strip as usual. * `torch._sparse_semi_structured_apply` This function takes in an input tensor and thread masks from the above function and returns a packed representation and metadata from applying thread masks to the input tensor. * `torch._sparse_semi_structured_apply_dense` This function does the same thing as above but instead of returning the tensor in the sparse representation it returns it in the dense representation The subclasses have also been updated to add a new `prune_dense_static_sort` classmethod to create sparse tensors with this format. I've added some additional documentatino on how to calculate the compressed tensors needed to create a SparseSemiStructuredTensor oneself. To this end, there are two new helper functions added: `sparse_semi_structured_tile` `compute_compressed_swizzled_bitmask` cc albanD Differential Revision: [D56190801](https://our.internmc.facebook.com/intern/diff/D56190801) [ghstack-poisoned]
diff --git a/test/test_sparse_semi_structured.py b/test/test_sparse_semi_structured.py
@@ -564,6 +564,8 @@ class TestSparseSemiStructuredTraining(TestCase):
     def setUp(self):
         if not _IS_SM8X:
             self.skipTest('Only runs on SM80')
+        if IS_WINDOWS:
+            self.skipTest('CUTLASS not supported on windows')
 
 
     @training_dtypes
