Make Python interface work for non-SM80 targets (#726)

* Make Python interface work for non-SM80 targets

* Remove line in README

tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py

@@ -2,9 +2,11 @@
 from pycutlass.conv2d_operation import *
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
 
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dDgradImplicitGemmF16nhwcF16nhwcF16nhwcTensorOpF16SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Dgrad_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dDgradImplicitGemmF16nhwcF16nhwcF32nhwcTensorOpF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Dgrad_Optimized_ImplicitGemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_unity_stride_stage3(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py

@@ -3,8 +3,11 @@ import pycutlass
 from pycutlass.conv2d_operation import *
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dDgradImplicitGemmF32nhwcF32nhwcF32nhwcSimtF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Fprop_Analytic_ImplicitGemm_f32nhwc_f32nhwc_f32nhwc_simt_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dDgradImplicitGemmTF32nhwcTF32nhwcTF32nhwcTensorOpF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Dgrad_Analytic_ImplicitGemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py

@@ -1,8 +1,11 @@
 # test/unit/conv/device/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu
 import pycutlass
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 def conv2d_few_channel_problemsizes(channels):
     problem_sizes = [
         cutlass.conv.Conv2dProblemSize(

tools/library/scripts/pycutlass/test/conv/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py

@@ -1,8 +1,11 @@
 # test/unit/conv/device/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.cu
 import pycutlass
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 def conv2d_fixed_channel_problemsizes(channels):
     problem_sizes = [
         cutlass.conv.Conv2dProblemSize(

tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dFpropImplicitGemmF16nhwcF16nhwcF16nhwcTensorOpF16SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Fprop_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dFpropImplicitGemmF16nhwcF16nhwcF32nhwcTensorOpF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Fprop_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py

@@ -3,8 +3,11 @@ import pycutlass
 from pycutlass.conv2d_operation import *
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dFpropImplicitGemmF32nhwcF32nhwcF32nhwcSimtF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Fprop_Analytic_ImplicitGemm_f32nhwc_f32nhwc_f32nhwc_simt_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dFpropImplicitGemmTF32nhwcTF32nhwcTF32nhwcTensorOpF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Fprop_Analytic_ImplicitGemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_strided_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dStridedDgradImplicitGemmF16NHWCF16NHWCF32NHWCTensorOpF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Strided_Dgrad_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_128x128_32x3_64x64x32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dWgradImplicitGemmF16nhwcF16nhwcF16nhwcTensorOpF16SM80(unittest.TestCase):
     def test_Device_Conv2d_Wgrad_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dWgradImplicitGemmF16nhwcF16nhwcF32nhwcTensorOpF32SM80(unittest.TestCase):
     def test_Device_Conv2d_Wgrad_Analytic_ImplicitGemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py

@@ -3,8 +3,11 @@ import pycutlass
 from pycutlass.conv2d_operation import *
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dWgradImplicitGemmF32nhwcF32nhwcF32nhwcSimtF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Wgrad_Analytic_ImplicitGemm_f32nhwc_f32nhwc_f32nhwc_simt_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py

@@ -2,8 +2,11 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class Conv2dWgradImplicitGemmTF32nhwcTF32nhwcTF32nhwcTensorOpF32SM80(unittest.TestCase):
     def test_SM80_Device_Conv2d_Wgrad_Optimized_ImplicitGemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/frontend/test_frontend.py

@@ -33,6 +33,7 @@
 import pycutlass
 import unittest
 from pycutlass import *
+from pycutlass.utils.device import device_cc
 import torch
 import cupy as cp
 
@@ -42,13 +43,18 @@ class Test_Frontend(unittest.TestCase):
         #
         # define the cutlass operator
         #
+        cc = device_cc()
         math_inst = MathInstruction(
             [1, 1, 1], cutlass.float32, cutlass.float32, cutlass.float32,
             cutlass.OpClass.Simt, MathOperation.multiply_add
         )
 
+        # Stages > 2 is supported only for compute capability 80 and beyond
+        stages = 4 if cc >= 80 else 2
+
+
         tile_description = TileDescription(
-            [128, 128, 8], 4, [2, 4, 1],
+            [128, 128, 8], stages, [2, 4, 1],
             math_inst
         )
 
@@ -69,7 +75,7 @@ class Test_Frontend(unittest.TestCase):
             math_inst.element_accumulator, cutlass.float32)
 
         self.operation = GemmOperationUniversal(
-            arch=80, tile_description=tile_description,
+            arch=cc, tile_description=tile_description,
             A=A, B=B, C=C, 
             epilogue_functor=epilogue_functor, 
             swizzling_functor=cutlass.IdentitySwizzle1

tools/library/scripts/pycutlass/test/gemm/gemm_bf16_sm80.py

@@ -4,7 +4,10 @@ from pycutlass.test import *
 import unittest
 
 from pycutlass.test.gemm_testbed import test_all_gemm
+from pycutlass.utils.device import device_cc
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class GemmBF16TensorOpSm80(unittest.TestCase):
     def SM80_Device_Gemm_bf16n_bf16n_f32t_tensor_op_f32_64x128x64_32x64x64(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/gemm/gemm_f16_sm80.py

@@ -4,8 +4,10 @@ from pycutlass.test import *
 import unittest
 
 from pycutlass.test.gemm_testbed import test_all_gemm
+from pycutlass.utils.device import device_cc
 
 
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class GemmF16Sm80(unittest.TestCase):
     def test_SM80_Device_Gemm_f32t_f32n_f32t_tensor_op_bf16_f32_128x128x32_64x64x32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/gemm/gemm_f32_sm80.py

@@ -5,8 +5,10 @@ from pycutlass.test import *
 import unittest
 
 from pycutlass.test.gemm_testbed import test_all_gemm
+from pycutlass.utils.device import device_cc
 
 
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
     def test_SM80_Device_Gemm_f32t_f32n_f32t_tensor_op_bf16_f32_128x128x32_64x64x32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/gemm/gemm_f64_sm80.py

@@ -4,7 +4,10 @@ from pycutlass.test import *
 import unittest
 
 from pycutlass.test.gemm_testbed import test_all_gemm
+from pycutlass.utils.device import device_cc
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class GemmF64TensorOpSm80(unittest.TestCase):
     def test_SM80_Device_Gemm_f64n_f64t_f64t_tensor_op_f64_32x32x16_16x16x16(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/gemm/gemm_grouped_sm80.py

@@ -4,8 +4,10 @@ from pycutlass.test import *
 import unittest
 
 from pycutlass.test.gemm_grouped_testbed import TestbedGrouped
+from pycutlass.utils.device import device_cc
 
 
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class GemmGroupedSm80(unittest.TestCase):
     def test_SM80_Device_GemmGrouped_f16n_f16t_f32n_tensor_op_f32_128x128x32_64x64x32(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/gemm/gemm_s8_sm80.py

@@ -5,7 +5,10 @@ from pycutlass.test import *
 import unittest
 
 from pycutlass.test.gemm_testbed import test_all_gemm
+from pycutlass.utils.device import device_cc
 
+
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 class GemmS8TensorOpF32Sm80(unittest.TestCase):
     def test_SM80_Device_Gemm_s8t_s8n_s8t_tensor_op_s32_64x64x64_32x32x64(self):
         math_inst = MathInstruction(

tools/library/scripts/pycutlass/test/unit/test_sm80.py

@@ -35,12 +35,14 @@
 import pycutlass
 from pycutlass import *
 from pycutlass.test import *
+from pycutlass.utils.device import device_cc
 import unittest
 
 #
 # Create GEMM operation
 #
 
+@unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
 def TestGemmOperator(gemm_kind, math_inst, layout, alignment, tiling, arch, mixed=False,
     epilogue_functor=None, swizzling_functor=cutlass.IdentitySwizzle1, **kwargs):
     """