setup.py

# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
#
# 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.
import os
from setuptools import setup, Extension, find_packages
import subprocess

from torch.utils.cpp_extension import BuildExtension, CUDAExtension, ROCM_HOME

from typing import Optional, Union
import subprocess
from pathlib import Path


version_dependent_macros = [
    '-DVERSION_GE_1_1',
    '-DVERSION_GE_1_3',
    '-DVERSION_GE_1_5',
]

extra_cuda_flags = [
    '-std=c++14',
    #'-maxrregcount=50',
    #'-U__CUDA_NO_HALF_OPERATORS__',
    #'-U__CUDA_NO_HALF_CONVERSIONS__',
    '--expt-relaxed-constexpr',
    '--expt-extended-lambda'
]

cc_flag = ['-DAMDGPU_TARGETS', 'gfx900;gfx906']

extra_cuda_flags += cc_flag


def get_sha(pytorch_root: Union[str, Path]) -> str:
    try:
        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=pytorch_root).decode('ascii').strip()
    except Exception:
        return 'Unknown'


def get_version_add(sha: Optional[str] = None) -> str:
    add_version_path = "version.py"
    if sha != 'Unknown':
        if sha is None:
            sha_path = os.getenv('OPENFOLD_DOWNLOAD_PATH', "")
            sha = get_sha(sha_path)
        version = 'git' + sha[:7]

    if os.getenv('OPENFOLD_BUILD_VERSION'):
        version_dtk = os.getenv('OPENFOLD_BUILD_VERSION', "")
        version += "." + version_dtk

    with open(add_version_path, encoding="utf-8", mode="w") as file:
        file.write("__version__='1.0.1'+'+{}'\n".format(version))
    file.close()


def get_version():
    get_version_add()
    version_file = 'version.py'
    with open(version_file, encoding='utf-8') as f:
        exec(compile(f.read(), version_file, 'exec'))
    return locals()['__version__']

setup(
    name='openfold',
    #version='1.0.1',
    version=get_version,
    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=False,
    package_data={
        "openfold": ['utils/kernel/csrc/*'],
        "": ["resources/stereo_chemical_props.txt"]
    },
    ext_modules=[CUDAExtension(
        name="attn_core_inplace_cuda",
        sources=[
            "openfold/utils/kernel/csrc/softmax_cuda.cpp",
            "openfold/utils/kernel/csrc/softmax_cuda_kernel.cu",
        ],
        include_dirs=[
            os.path.join(
                os.path.dirname(os.path.abspath(__file__)),
                'openfold/utils/kernel/csrc/'
            )
        ],
        extra_compile_args={
            'cxx': ['-O3'] + version_dependent_macros,
            'hipcc': (
                ['-O3', '--use_fast_math'] +
                version_dependent_macros +
                extra_cuda_flags
            ),
        }
    )],
    cmdclass={'build_ext': BuildExtension},
    classifiers=[
        'License :: OSI Approved :: Apache Software License',
        'Operating System :: POSIX :: Linux',
        'Programming Language :: Python :: 3.7,'
        'Topic :: Scientific/Engineering :: Artificial Intelligence',
    ],
)