81 lines
2.4 KiB
Python
81 lines
2.4 KiB
Python
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",
|
|
)
|