Update on "[sparse] semi-structured sparse + torch.compile support"

jcaip · jcaip · commit 575d075322f6 · 2023-10-23T13:46:11.000-07:00
Summary:

This PR adds in torch.compile support for semi-structured sparsity,
using the subclass tracing bdhirsh added.

Based on wether we are using cuSPARSELt or CUTLASS, we return a
different representation of the inner tensors.

Test Plan:
```
python test/test_sparse_semi_structured.py -k compile
```

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
diff --git a/test/test_sparse_semi_structured.py b/test/test_sparse_semi_structured.py
@@ -127,7 +127,6 @@ def setUp(self):
     def tearDown(self):
         super().tearDown()
 
-
     @unittest.skipIf(IS_WINDOWS, "torch.compile not support on windows")
     def test_mlp_contiguous_relu_compile(self):
         """
@@ -148,8 +147,8 @@ def forward(self, x):
                 return torch.nn.functional.relu(x)
 
         def _test_mlp_contiguous_relu_compile(backend, dense_input_shape):
-            SparseSemiStructuredTensor._FORCE_CUTLASS = False
-            input = torch.rand((128, 128), device="cuda").half()
+            SparseSemiStructuredTensor._FORCE_CUTLASS = backend == "cutlass"
+            input = torch.rand(dense_input_shape, device="cuda").half()
 
             model = Model().eval().cuda().half()
             mod_linear = model.linear
@@ -175,11 +174,6 @@ class TestSparseSemiStructured(TestCase):
     def setUp(self):
         if not _IS_SM8X:
             self.skipTest('Only runs on SM80')
-        super().setUp()
-
-    def tearDown(self):
-        super().tearDown()
-
 
     @dtypes(*SEMI_STRUCTURED_SUPPORTED_DTYPES)
     @parametrize("backend", SEMI_STRUCTURED_SUPPORTED_BACKENDS)
