4747
4848_IS_SM8X = False
4949_IS_SM9X = False
50+ _IS_HIPSPARSELT_AVAILABLE = False
5051
5152if torch .cuda .is_available ():
5253 _IS_SM8X = torch .cuda .get_device_capability (0 )[0 ] == 8
5354 _IS_SM9X = torch .cuda .get_device_capability (0 )[0 ] == 9
54-
55+ _IS_HIPSPARSELT_AVAILABLE = torch . version . hip is not None and tuple ( int ( v ) for v in torch . version . hip . split ( '.' )[: 2 ]) > ( 6 , 4 )
5556 # CUTLASS kernels only work for Ampere
5657 if _IS_SM8X :
5758 SEMI_STRUCTURED_SUPPORTED_BACKENDS ["cutlass" ] = SparseSemiStructuredTensorCUTLASS
5859
5960 # add cuSPASRELt tests if available
60- if torch .backends .cusparselt .is_available () and (_IS_SM8X or _IS_SM9X ):
61+ if torch .backends .cusparselt .is_available () and (_IS_SM8X or _IS_SM9X or _IS_HIPSPARSELT_AVAILABLE ):
6162 SEMI_STRUCTURED_SUPPORTED_BACKENDS ["cusparselt" ] = SparseSemiStructuredTensorCUSPARSELT
6263
6364inference_dtypes = dtypes (torch .float16 , torch .bfloat16 , torch .int8 )
@@ -223,6 +224,7 @@ def forward(self, x):
223224
224225 @unittest .skipIf (IS_WINDOWS , "torch.compile not supported on windows" )
225226 @unittest .skipIf ("cusparselt" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS , "cusparselt not supported on this machine" )
227+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
226228 def test_mlp_contiguous_relu_compile_cusparselt (self ):
227229 """
228230 test for cuSPASRELt meta registrations (_cslt_sparse_mm) + torch.compile
@@ -233,6 +235,7 @@ def test_mlp_contiguous_relu_compile_cusparselt(self):
233235
234236 @unittest .skipIf ("cutlass" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS , "cutlass not supported on this machine" )
235237 @unittest .skipIf (IS_WINDOWS , "torch.compile not supported on windows" )
238+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
236239 def test_mlp_contiguous_relu_compile_cutlass (self ):
237240 """
238241 test for CUTLASS meta registrations (_sparse_semi_structured_addmm) + torch.compile
@@ -243,6 +246,7 @@ def test_mlp_contiguous_relu_compile_cutlass(self):
243246
244247 @unittest .skipIf (IS_WINDOWS , "torch.compile not supported on windows" )
245248 @unittest .skipIf ("cusparselt" not in SEMI_STRUCTURED_SUPPORTED_BACKENDS , "cusparselt not supported on this machine" )
249+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
246250 def test_sp24_compile (self ) -> None :
247251 x = torch .randn ([1024 , 512 ], device = "cuda" , dtype = torch .float16 , requires_grad = True )
248252
@@ -571,6 +575,7 @@ def setUp(self):
571575
572576
573577 @training_dtypes
578+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
574579 def test_prune_dense_static_sort (self , dtype ) -> None :
575580 # Ideally we would like to clone and compare, but that won't work because the sorting order will be different
576581 # instead we pass the pruned matrix to the CUDA implementation and preserve the sparsity pattern.
@@ -615,6 +620,7 @@ def test_prune_dense_static_sort(self, dtype) -> None:
615620
616621 @training_dtypes
617622 @parametrize_backends
623+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
618624 def test_pruning_algo_largest_abs_values_greedy (self , dtype , backend ) -> None :
619625 inp = torch .tensor (
620626 [[4 , 3 , 2 , 1 ], [- 1 , - 3 , 0.6 , 0.5 ], [1 , 2 , 3 , 4 ], [10 , 2 , - 1 , 5 ]],
@@ -651,6 +657,7 @@ def test_gemm(self, dtype) -> None:
651657
652658 @training_dtypes
653659 @parametrize_backends
660+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
654661 def test_pack_both_ways_meta_correctness (self , dtype , backend ) -> None :
655662 M , N = 128 , 256
656663 # Construct x to make sure we always have exactly 8 elements per 4x4 tile
@@ -684,6 +691,7 @@ def test_pack_both_ways_meta_correctness(self, dtype, backend) -> None:
684691 torch .testing .assert_close (ref_gemm , pack_gemm , ** atol_rtol_kw [dtype ])
685692
686693 @training_dtypes
694+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
687695 def test_pack_both_ways_id (self , dtype ) -> None :
688696 N = 512
689697 torch .manual_seed (0 )
@@ -718,6 +726,7 @@ def test_pack_both_ways_id(self, dtype) -> None:
718726 ), f"packed_t is wrong at pos: ({ max_diff // N } { max_diff % N }
719727
720728 @training_dtypes
729+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
721730 def test_pack_both_ways_edge_case1 (self , dtype ) -> None :
722731 # In this case, the heuristic will keep 7 values out of 16
723732 # instead of 8. let's see how the kernel handles this
@@ -742,6 +751,7 @@ def test_pack_both_ways_edge_case1(self, dtype) -> None:
742751 assert packed_t [0 , 1 ].item () == 0
743752
744753 @training_dtypes
754+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
745755 def test_sp24_apply (self , dtype ) -> None :
746756 M , N = 256 , 1024
747757 x = torch .randn ([M , N ], dtype = dtype , device = "cuda" )
@@ -757,6 +767,7 @@ def test_sp24_apply(self, dtype) -> None:
757767 torch .testing .assert_close (packed_t , packed_t2 )
758768
759769 @training_dtypes
770+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
760771 def test_sp24_apply_dense (self , dtype ) -> None :
761772 M , N = 256 , 1024
762773 x = torch .randn ([M , N ], dtype = dtype , device = "cuda" )
@@ -794,6 +805,7 @@ def test_sp24_apply_dense(self, dtype) -> None:
794805
795806
796807 @training_dtypes
808+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
797809 def test_sp24_matmuls (self , dtype ) -> None :
798810 M , N , K = 64 , 256 , 1024
799811 a = torch .randn ([M , K ], device = "cuda" , dtype = dtype )
@@ -828,6 +840,7 @@ def test_sp24_matmuls(self, dtype) -> None:
828840 a_s .t () @ a , (a * a_m ).t () @ a , rtol = 1e-1 , atol = 1e-1
829841 )
830842
843+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
831844 def test_sp24_matmuls_mat_vec (self ) -> None :
832845 a = torch .randn ([64 , 128 ], device = "cuda" , dtype = torch .float16 )
833846 b = torch .randn ([128 ], device = "cuda" , dtype = torch .float16 )
@@ -837,7 +850,7 @@ def test_sp24_matmuls_mat_vec(self) -> None:
837850 with pytest .raises (NotImplementedError ):
838851 torch .testing .assert_close (a_s @ b , (a * a_m ) @ b , ** atol_rtol_kw [a .dtype ])
839852
840-
853+ @ unittest . skipIf ( TEST_WITH_ROCM , "Not supported on ROCm" )
841854 def test_sp24_matmuls_bmm (self ) -> None :
842855 a = torch .randn ([64 , 128 ], device = "cuda" , dtype = torch .float16 )
843856 b = torch .randn ([5 , 6 , 128 ], device = "cuda" , dtype = torch .float16 )
@@ -988,6 +1001,7 @@ def run_test(m, n, k, device, dtype, dtype_out, use_input, rtol, atol):
9881001
9891002 @unittest .skipIf (not HAS_GPU , "Inductor+gpu needs triton and recent GPU arch" )
9901003 @inference_dtypes
1004+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
9911005 def test_conversions (self , device , dtype ):
9921006
9931007 def run_test (r , c , device , dtype ):
@@ -1016,6 +1030,7 @@ def run_test(r, c, device, dtype):
10161030
10171031 @unittest .skipIf (not HAS_GPU , "Inductor+gpu needs triton and recent GPU arch" )
10181032 @inference_dtypes
1033+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
10191034 def test_conversions_all_patterns (self , device , dtype ):
10201035 r , c = 32 , 128
10211036
@@ -1135,6 +1150,7 @@ def test_cslt_sparse_mm_mixed_dtype(self, dense_input_shape, out_dtype, device):
11351150
11361151 @unittest .skip ("cuSPARSELt v0.6.x does not support bfloat/float16 alpha scaling" )
11371152 @training_dtypes
1153+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
11381154 def test_cslt_sparse_mm_alpha (self , dtype , device ):
11391155 A = torch .Tensor ([0 , 0 , 1 , 1 ]).tile ((128 , 64 )).to (dtype ).cuda ()
11401156 B = torch .ones ((256 , 128 ), device = device ).to (dtype )
@@ -1151,6 +1167,7 @@ def test_cslt_sparse_mm_alpha(self, dtype, device):
11511167 torch .testing .assert_close (sparse_result , dense_result , rtol = 1e-3 , atol = 1e-3 )
11521168
11531169 @parametrize ("out_dtype" , [torch .float16 , torch .bfloat16 , torch .int32 ])
1170+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
11541171 def test_cslt_sparse_mm_alpha_compile_autotune (self , device , out_dtype ):
11551172 A = torch .Tensor ([0 , 0 , 1 , 1 ]).tile ((128 , 64 )).to (torch .int8 ).to (device )
11561173 B = torch .ones ((128 , 256 ), device = device , dtype = torch .int8 ).t ()
@@ -1172,6 +1189,7 @@ def get_dense_result():
11721189 torch .testing .assert_close (sparse_result .cpu (), get_dense_result (), rtol = 1e-3 , atol = 1e-3 )
11731190
11741191 @parametrize ("out_dtype" , [torch .float16 , torch .bfloat16 , torch .int32 ])
1192+ @unittest .skipIf (TEST_WITH_ROCM , "Not supported on ROCm" )
11751193 def test_cslt_sparse_mm_alpha_mixed_dtype (self , out_dtype , device ):
11761194 A = torch .Tensor ([0 , 0 , 10 , 10 ]).tile ((128 , 64 )).to (torch .int8 ).cuda ()
11771195 B = torch .ones ((128 , 256 ), device = device ).to (torch .int8 ).t ()
0 commit comments