CUTLASS 2.10 (#615)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
2022-09-03 18:48:46 -04:00
parent ca23ff7924
commit b72cbf957d
289 changed files with 43708 additions and 2513 deletions
@@ -0,0 +1,79 @@
+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"])
+    ]
+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'
+    ],
+    cmdclass={
+        'rmm': BuildRMM
+    },
+    ext_modules=ext_modules,
+    python_requires=">=3.6",
+)