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/README.md

@@ -102,8 +102,10 @@ Examples can be found in [$CUTLASS_PATH/examples/40_cutlass_py](examples/40_cutl
 ## Test
 The test cases are listed in `$CUTLASS_PATH//tools/library/scripts/pycutlass/test`. The unit test can be run with
 ```shell
+# Each of these tests are only supported on devices with compute capability of SM80. For other devices,
+# see the basic examples in $CUTLASS_PATH/examples/40_cutlass_py
 cd $CUTLASS_PATH/tools/library/scripts/pycutlass/test/unit && python test_sm80.py
-cd $CUTLASS_PATH/tools/library/scripts/pycutlass/test/example && run_all_example.sh
+cd $CUTLASS_PATH/tools/library/scripts/pycutlass/test/example && bash run_all_example.sh
 ```
 
 ## build documentation

tools/library/scripts/pycutlass/src/pycutlass/compiler.py

@@ -308,7 +308,7 @@ class ArtifactManager:
         cmd = "echo '%s'|g++ -x c++ -fpermissive -w -fPIC" % source_buffer_host
         for opt in options:
             opt = opt.decode("utf-8")
-            if opt not in ['-default-device', '-std=c++11', '-arch=sm_80', '-Xcicc', '-Xllc']:
+            if opt not in ['-default-device', '-std=c++11', '-Xcicc', '-Xllc'] and '-arch=sm_' not in opt:
                 if '--include-path=' in opt:
                     cmd += " " + opt.replace('--include-path=', '-I')
                 else:

tools/library/scripts/pycutlass/src/pycutlass/utils/device.py

@@ -0,0 +1,76 @@
+#################################################################################################
+#
+# Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+#################################################################################################
+
+"""
+Utility functions for interacting with the device
+"""
+
+from cuda import cudart
+
+
+def check_cuda_errors(result: list):
+    """
+    Checks whether `result` contains a CUDA error raises the error as an exception, if so. Otherwise,
+    returns the result contained in the remaining fields of `result`.
+
+    :param result: the results of the `cudart` method, consisting of an error code and any method results
+    :type result: list
+
+    :return: non-error-code results from the `results` parameter
+    """
+    # `result` is of the format : (cudaError_t, result...)
+    err = result[0]
+    if err.value:
+        raise RuntimeError("CUDA error: {}".format(cudart.cudaGetErrorName(err)))
+
+    if len(result) == 1:
+        return None
+    elif len(result) == 2:
+        return result[1]
+    else:
+        return result[1:]
+
+
+def device_cc(device: int = 0) -> int:
+    """
+    Returns the compute capability of the device with ID `device`.
+
+    :param device: ID of the device to query
+    :type device: int
+
+    :return: compute capability of the queried device (e.g., 80 for SM80)
+    :rtype: int
+    """
+    deviceProp = check_cuda_errors(cudart.cudaGetDeviceProperties(device))
+    major = str(deviceProp.major)
+    minor = str(deviceProp.minor)
+    return int(major + minor)

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):
     """