Fix kernel installation bug

scripts/utils.py

 import argparse
+import ctypes
 from datetime import date
+import sys
 
 
 def add_data_args(parser: argparse.ArgumentParser):
@@ -31,7 +33,7 @@ def add_data_args(parser: argparse.ArgumentParser):
         '--kalign_binary_path', type=str, default='/usr/bin/kalign'
     )
     parser.add_argument(
-        '--max_template_date', type=str, 
+        '--max_template_date', type=str,
         default=date.today().strftime("%Y-%m-%d"),
     )
     parser.add_argument(
@@ -40,3 +42,59 @@ def add_data_args(parser: argparse.ArgumentParser):
     parser.add_argument(
         '--release_dates_path', type=str, default=None
     )
+
+
+def get_nvidia_cc():
+    """
+    Returns a tuple containing the Compute Capability of the first GPU
+    installed in the system (formatted as a tuple of strings) and an error
+    message. When the former is provided, the latter is None, and vice versa.
+
+    Adapted from script by Jan Schlüte a
+    https://gist.github.com/f0k/63a664160d016a491b2cbea15913d549
+    """
+    CUDA_SUCCESS = 0
+
+    libnames = ('libcuda.so', 'libcuda.dylib', 'cuda.dll')
+    for libname in libnames:
+        try:
+            cuda = ctypes.CDLL(libname)
+        except OSError:
+            continue
+        else:
+            break
+    else:
+        raise OSError("could not load any of: " + ' '.join(libnames))
+
+    nGpus = ctypes.c_int()
+    cc_major = ctypes.c_int()
+    cc_minor = ctypes.c_int()
+
+    result = ctypes.c_int()
+    device = ctypes.c_int()
+    error_str = ctypes.c_char_p()
+
+    result = cuda.cuInit(0)
+    if result != CUDA_SUCCESS:
+        err = cuda.cuGetErrorString(result, ctypes.byref(error_str))
+        return None, err.value.decode()
+    result = cuda.cuDeviceGetCount(ctypes.byref(nGpus))
+    if result != CUDA_SUCCESS:
+        err = cuda.cuGetErrorString(result, ctypes.byref(error_str))
+        return None, err.value.decode()
+
+    if(nGpus.value < 1):
+        return None, err.value.decode()
+
+    result = cuda.cuDeviceGet(ctypes.byref(device), 0)
+    if result != CUDA_SUCCESS:
+        err = cuda.cuGetErrorString(result, ctypes.byref(error_str))
+        return None, err.value.decode()
+
+    if cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), device) != CUDA_SUCCESS:
+        return None, "Compute Capability not found"
+
+    major = cc_major.value
+    minor = cc_minor.value
+
+    return (major, minor), None

setup.py

@@ -18,6 +18,8 @@ import subprocess
 
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
 
+from scripts.utils import get_nvidia_cc
+
 
 version_dependent_macros = [
     '-DVERSION_GE_1_1',
@@ -26,11 +28,11 @@ version_dependent_macros = [
 ]
 
 extra_cuda_flags = [
-    '-std=c++14', 
-    '-maxrregcount=50', 
+    '-std=c++14',
+    '-maxrregcount=50',
     '-U__CUDA_NO_HALF_OPERATORS__',
-    '-U__CUDA_NO_HALF_CONVERSIONS__', 
-    '--expt-relaxed-constexpr', 
+    '-U__CUDA_NO_HALF_CONVERSIONS__',
+    '--expt-relaxed-constexpr',
     '--expt-extended-lambda'
 ]
 
@@ -44,11 +46,24 @@ def get_cuda_bare_metal_version(cuda_dir):
 
     return raw_output, bare_metal_major, bare_metal_minor
 
-cc_flag = ['-gencode', 'arch=compute_70,code=sm_70']
+compute_capabilities = set()
+compute_capabilities.add((7, 0))
 _, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
 if int(bare_metal_major) >= 11:
-    cc_flag.append('-gencode')
-    cc_flag.append('arch=compute_80,code=sm_80')
+    compute_capabilities.add((8, 0))
+
+compute_capability, _ = get_nvidia_cc()
+if compute_capability is not None:
+    compute_capabilities.add(compute_capability)
+
+print(compute_capabilities)
+
+cc_flag = []
+for major, minor in list(compute_capabilities):
+    cc_flag.extend([
+        '-gencode',
+        f'arch=compute_{major}{minor},code=sm_{major}{minor}',
+    ])
 
 extra_cuda_flags += cc_flag
 
@@ -75,15 +90,15 @@ setup(
         ],
         include_dirs=[
             os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), 
+                os.path.dirname(os.path.abspath(__file__)),
                 'openfold/utils/kernel/csrc/'
             )
         ],
         extra_compile_args={
             'cxx': ['-O3'] + version_dependent_macros,
             'nvcc': (
-                ['-O3', '--use_fast_math'] + 
-                version_dependent_macros + 
+                ['-O3', '--use_fast_math'] +
+                version_dependent_macros +
                 extra_cuda_flags
             ),
         }
@@ -92,7 +107,10 @@ setup(
     classifiers=[
         'License :: OSI Approved :: Apache Software License',
         'Operating System :: POSIX :: Linux',
-        'Programming Language :: Python :: 3.7,' 
+        'Programming Language :: Python :: 3.7,'
         'Topic :: Scientific/Engineering :: Artificial Intelligence',
     ],
 )
+
+
+