NvidiaWarp-GarmentCode/warp/native/cutlass/tools/library/scripts/pycutlass/setup.py

import distutils.cmd
from setuptools import setup
import setuptools.command.build_py
import os

# build rmm dependency
class BuildRMM(distutils.cmd.Command):
    user_options = []
    def initialize_options(self):
        pass
    def finalize_options(self):
        pass
    def run(self):
        try:
            import rmm
        except ImportError:
            print("installing rmm")
            os.system("git clone -b branch-22.08 --recurse-submodules https://github.com/rapidsai/rmm.git")
            os.chdir("./rmm")
            os.system("./build.sh librmm rmm")
            os.chdir("./python")
            os.system("python setup.py build_ext --inplace")
            os.system("python setup.py install")

cutlass_path = os.getenv('CUTLASS_PATH')
assert cutlass_path is not None, "Environment variable 'CUTLASS_PATH' is not defined."
cuda_install_path = os.getenv('CUDA_INSTALL_PATH')
assert cuda_install_path is not None, "Environment variable 'CUDA_INSTALL_PATH' is not defined."

ext_modules = []

try:
    from pybind11.setup_helpers import Pybind11Extension, build_ext
    include_dirs = [
        cutlass_path + "/include",
        cuda_install_path + "/include",
        cutlass_path + "/tools/util/include",
        cutlass_path + "/test",
        cutlass_path + "/tools/library/scripts/pycutlass/googletest/googletest/include"
    ]

    ext_modules = [
        Pybind11Extension("cutlass",
                          ["src/cpp/cutlass.cpp"],
                          include_dirs=include_dirs,
                          extra_compile_args=["-fpermissive", "-w"])
    ]
except ImportError:
    pass

setup(
    name="PyCutlass",
    version="0.0.1",
    author="Zhaodong Chen; Andrew Kerr; Haicheng Wu; Szymon Migacz; Graham Markall",
    author_email="zhaodongc@nvidia.com",
    description="Python interface for CUTLASS",
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    package_dir={"": "src"},
    packages=['pycutlass', 'pycutlass.utils', 'pycutlass.test'],
    setup_requires=["pybind11", "numpy<1.23"],
    install_requires=[
        "numpy<1.23",
        'pybind11',
        'cuda-python<11.7.0',
        'typeguard',
        'bfloat16',
        'typing',
        'scikit-build',
        'treelib'
    ],
    cmdclass={
        'rmm': BuildRMM
    },
    ext_modules=ext_modules,
    python_requires=">=3.6",
)