Merge pull request #252 from p-durandin/fold-cpu

Fix for CPU only install

openfold/utils/kernel/csrc/softmax_cuda_stub.cpp

+// Copyright 2021 AlQuraishi Laboratory
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// modified from fastfold/model/fastnn/kernel/cuda_native/csrc/softmax_cuda.cpp
+
+#include <torch/extension.h>
+
+void attn_softmax_inplace_forward_(
+    at::Tensor input, 
+    long long rows, int cols
+)
+{
+    throw std::runtime_error("attn_softmax_inplace_forward_ not implemented on CPU");
+};
+void attn_softmax_inplace_backward_(
+    at::Tensor output, 
+    at::Tensor d_ov,
+    at::Tensor values,
+    long long rows, 
+    int cols_output,
+    int cols_values
+)
+{
+    throw std::runtime_error("attn_softmax_inplace_backward_ not implemented on CPU");
+};
\ No newline at end of file

setup.py

@@ -16,7 +16,7 @@ import os
 from setuptools import setup, Extension, find_packages
 import subprocess
 
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
+from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME
 
 from scripts.utils import get_nvidia_cc
 
@@ -37,14 +37,18 @@ extra_cuda_flags = [
 ]
 
 def get_cuda_bare_metal_version(cuda_dir):
-    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
-    output = raw_output.split()
-    release_idx = output.index("release") + 1
-    release = output[release_idx].split(".")
-    bare_metal_major = release[0]
-    bare_metal_minor = release[1][0]
-
-    return raw_output, bare_metal_major, bare_metal_minor
+    if cuda_dir==None:
+        print("CUDA is not found, cpu version is installed")
+        return None, -1, 0
+    else:
+        raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
+        output = raw_output.split()
+        release_idx = output.index("release") + 1
+        release = output[release_idx].split(".")
+        bare_metal_major = release[0]
+        bare_metal_minor = release[1][0]
+        
+        return raw_output, bare_metal_major, bare_metal_minor
 
 compute_capabilities = set([
     (3, 7), # K80, e.g.
@@ -70,22 +74,8 @@ for major, minor in list(compute_capabilities):
 
 extra_cuda_flags += cc_flag
 
-
-setup(
-    name='openfold',
-    version='1.0.0',
-    description='A PyTorch reimplementation of DeepMind\'s AlphaFold 2',
-    author='Gustaf Ahdritz & DeepMind',
-    author_email='gahdritz@gmail.com',
-    license='Apache License, Version 2.0',
-    url='https://github.com/aqlaboratory/openfold',
-    packages=find_packages(exclude=["tests", "scripts"]),
-    include_package_data=True,
-    package_data={
-        "openfold": ['utils/kernel/csrc/*'],
-        "": ["resources/stereo_chemical_props.txt"]
-    },
-    ext_modules=[CUDAExtension(
+if bare_metal_major != -1:
+    modules = [CUDAExtension(
         name="attn_core_inplace_cuda",
         sources=[
             "openfold/utils/kernel/csrc/softmax_cuda.cpp",
@@ -105,7 +95,34 @@ setup(
                 extra_cuda_flags
             ),
         }
-    )],
+    )]
+else:
+    modules = [CppExtension(
+        name="attn_core_inplace_cuda",
+        sources=[
+            "openfold/utils/kernel/csrc/softmax_cuda.cpp",
+            "openfold/utils/kernel/csrc/softmax_cuda_stub.cpp",
+        ],
+        extra_compile_args={
+            'cxx': ['-O3'],
+        }
+    )]
+
+setup(
+    name='openfold',
+    version='1.0.0',
+    description='A PyTorch reimplementation of DeepMind\'s AlphaFold 2',
+    author='Gustaf Ahdritz & DeepMind',
+    author_email='gahdritz@gmail.com',
+    license='Apache License, Version 2.0',
+    url='https://github.com/aqlaboratory/openfold',
+    packages=find_packages(exclude=["tests", "scripts"]),
+    include_package_data=True,
+    package_data={
+        "openfold": ['utils/kernel/csrc/*'],
+        "": ["resources/stereo_chemical_props.txt"]
+    },
+    ext_modules=modules,
     cmdclass={'build_ext': BuildExtension},
     classifiers=[
         'License :: OSI Approved :: Apache Software License',
@@ -113,7 +130,4 @@ setup(
         'Programming Language :: Python :: 3.7,'
         'Topic :: Scientific/Engineering :: Artificial Intelligence',
     ],
-)
-
-
-
+)
\ No newline at end of file